ViT-Style Real-Data Example

Runnable torchlean vit example. It trains a compact ViT-style image classifier on a prepared CIFAR-10 minibatch: patch embedding by convolution, token reshape, transformer block, and linear head.

The reusable model wiring lives behind the public TorchLean.nn.models.ViT constructor. The command adds CIFAR loader construction and the step-limited training loop.

python3 scripts/datasets/download_example_data.py --cifar10
lake build -R -K cuda=true && lake exe torchlean vit --cuda --n-total 1 --steps 1

This command is a small runtime check. Larger image-token runs belong in runtime profiling work, not the default quick path:

lake build -R -K cuda=true
lake exe torchlean vit --cuda --fast-kernels --n-total 1 --steps 1

def NN.Examples.Models.Vision.Vit.exeName : String

CLI subcommand name used in terminal banners and parser errors.

def NN.Examples.Models.Vision.Vit.defaultLogJson : System.FilePath

Default JSON loss-curve path for this command.

def NN.Examples.Models.Vision.Vit.batch : ℕ

Static minibatch size for the ViT example.

The batch axis is part of the checked model type, so changing this value changes the input and output shapes at compile time.

def NN.Examples.Models.Vision.Vit.inC : ℕ

CIFAR image channels.

def NN.Examples.Models.Vision.Vit.inH : ℕ

Height of the CIFAR crop used by this runnable ViT command.

def NN.Examples.Models.Vision.Vit.inW : ℕ

Width of the CIFAR crop used by this runnable ViT command.

def NN.Examples.Models.Vision.Vit.patchH : ℕ

Patch height used by the convolutional patch embedding.

def NN.Examples.Models.Vision.Vit.patchW : ℕ

Patch width used by the convolutional patch embedding.

def NN.Examples.Models.Vision.Vit.stride : ℕ

Patch stride; equal to patch size here, so patches do not overlap.

def NN.Examples.Models.Vision.Vit.padding : ℕ

No zero-padding for the patch embedding.

def NN.Examples.Models.Vision.Vit.dModel : ℕ

Transformer feature width.

CIFAR rows are cropped before training. A 2×2 patch covers the whole crop here, so the command exercises the ViT path with one image token and a small classifier head.

def NN.Examples.Models.Vision.Vit.outDim : ℕ

CIFAR class count, hence the output-logit width.

def NN.Examples.Models.Vision.Vit.numHeads : ℕ

Number of attention heads in the single encoder block.

def NN.Examples.Models.Vision.Vit.headDim : ℕ

Per-head feature width; numHeads * headDim = dModel.

def NN.Examples.Models.Vision.Vit.ffnHidden : ℕ

Feed-forward hidden width inside the encoder block.

def NN.Examples.Models.Vision.Vit.cfg : TorchLean.nn.models.VitConfig

Shared ViT configuration used by shapes and the reusable public model constructor.

@[reducible, inline]

abbrev NN.Examples.Models.Vision.Vit.σ : Spec.Shape

@[reducible, inline]

abbrev NN.Examples.Models.Vision.Vit.τ : Spec.Shape

def NN.Examples.Models.Vision.Vit.model : TorchLean.nn.M (TorchLean.nn.Sequential σ τ)

Compact ViT-style classifier from the public model API.

The constructor builds patch embedding, token reshape, one encoder block, and the classifier head.

def NN.Examples.Models.Vision.Vit.train (opts : TorchLean.Options) (flags : RealData.CifarModelTrainFlags) : IO TorchLean.Trainer.TrainSummary

Train the CIFAR ViT with the public Trainer surface.

def NN.Examples.Models.Vision.Vit.main (args : List String) : IO UInt32

CLI entrypoint for CIFAR ViT training; CUDA is the maintained validation path.