Masked Autoencoder CIFAR Example

This is the compact ViT-MAE-style training path in TorchLean.

The data path is explicit:

1. load real CIFAR-10 .npy arrays through Data;
2. take a typed image batch with shape [batch, channels, height, width];
3. hide deterministic image patches with ssl.imagePatchMaeSample;
4. run a ViT encoder over patch tokens;
5. train a decoder head to reconstruct the original image vector.

The architecture uses one transformer encoder block and a linear pixel decoder rather than a large asymmetric MAE decoder. The important pieces are the MAE pieces exercised by the example: image patch masking, patch embedding, transformer tokens, and reconstruction of the original image.

def NN.Examples.Models.Generative.Mae.exeName : String

Command name used in error messages and CLI output.

def NN.Examples.Models.Generative.Mae.defaultLogJson : System.FilePath

Default JSON loss-curve path for this command.

def NN.Examples.Models.Generative.Mae.batch : ℕ

CIFAR minibatch size used by the typed MAE command.

def NN.Examples.Models.Generative.Mae.inC : ℕ

Number of CIFAR image channels.

def NN.Examples.Models.Generative.Mae.inH : ℕ

Cropped CIFAR image height for the compact runnable example.

def NN.Examples.Models.Generative.Mae.inW : ℕ

Cropped CIFAR image width for the compact runnable example.

def NN.Examples.Models.Generative.Mae.patchH : ℕ

Patch height for the image-to-token projection.

def NN.Examples.Models.Generative.Mae.patchW : ℕ

Patch width for the image-to-token projection.

def NN.Examples.Models.Generative.Mae.stride : ℕ

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

def NN.Examples.Models.Generative.Mae.padding : ℕ

Zero padding around the image before patch extraction.

def NN.Examples.Models.Generative.Mae.dModel : ℕ

Width of each patch token after projection into the encoder stream.

def NN.Examples.Models.Generative.Mae.numHeads : ℕ

Number of self-attention heads in the compact ViT encoder.

def NN.Examples.Models.Generative.Mae.headDim : ℕ

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

def NN.Examples.Models.Generative.Mae.ffnHidden : ℕ

Hidden width of the feed-forward block inside the encoder.

def NN.Examples.Models.Generative.Mae.reconDim : ℕ

Number of reconstructed flattened pixels predicted by the decoder head.

def NN.Examples.Models.Generative.Mae.cfg : TorchLean.nn.models.VitMaeConfig

Small ViT-MAE configuration.

The command crops CIFAR images to 2×2, uses one image patch, and reconstructs a tiny prefix of the flattened image. That keeps MAE in the runnable quick-check suite while still checking the patch masking, patch embedding, transformer token, decoder, data loading, and CUDA training path.

def NN.Examples.Models.Generative.Mae.maskPeriod : ℕ

Hide one patch-index class every four patch positions.

The image remains an image tensor; the mask zeros whole patch regions before patch embedding.

def NN.Examples.Models.Generative.Mae.maskOffset : ℕ

Phase of the deterministic patch mask. Changing this selects a different patch-index class.

@[reducible, inline]

abbrev NN.Examples.Models.Generative.Mae.σ : Shape

Input shape: a real batched CIFAR image tensor.

@[reducible, inline]

abbrev NN.Examples.Models.Generative.Mae.τ : Shape

Output shape: flattened image reconstruction.

def NN.Examples.Models.Generative.Mae.cifarClasses : ℕ

CIFAR-10 images are stored as 3 × 32 × 32 tensors.

def NN.Examples.Models.Generative.Mae.model : TorchLean.nn.M (TorchLean.nn.Sequential σ τ)

Construct the trainable model.

The architecture lives in the public self-supervised model API; this example only chooses a config, loads data, and trains it.

def NN.Examples.Models.Generative.Mae.mkMaeSample (b : TorchLean.SupervisedSample Float (TorchLean.Shape.images cfg.batch cfg.inC cfg.inH cfg.inW) (TorchLean.Shape.mat cfg.batch RealData.cifarClasses)) : TorchLean.SupervisedSample Float σ τ

Turn a typed CIFAR image batch into the compact MAE training sample.

The input stays an image tensor with some patches zeroed out. The target is the original image flattened to a vector because the current decoder head predicts a batched matrix.

def NN.Examples.Models.Generative.Mae.data (flags : RealData.CifarModelTrainFlags) : TorchLean.Trainer.Dataset σ τ

Public singleton dataset for masked-image reconstruction on one real CIFAR batch.

Like the compact vector generative examples, the sample itself is loaded as Float from the real data boundary, then cast into the runtime-selected scalar by the public dataset constructor.

def NN.Examples.Models.Generative.Mae.train (opts : TorchLean.Options) (flags : RealData.CifarModelTrainFlags) : IO (TorchLean.Trainer.TrainResult σ τ)

Train the compact MAE model with the public Trainer surface.

def NN.Examples.Models.Generative.Mae.main (args : List String) : IO UInt32

CLI entrypoint.

Useful flags:

• --cuda runs the public trainer on the CUDA runtime.
• --steps <n> controls optimization steps.
• --x <path> --y <path> selects custom CIFAR-style .npy arrays.
• --log <path> writes the standard TorchLean training log JSON.