Masked Autoencoder CIFAR Example

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

The data path is intentionally concrete:

1. load real CIFAR-10 .npy arrays through NN.API.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.

This is still intentionally small: 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 training curve location. data/ is intentionally ignored by git.

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

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

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

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

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

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

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

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

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

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

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

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

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

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

Small ViT-MAE configuration.

CIFAR-10 with 16×16 patches gives 2×2 = 4 patch tokens. The model embeds patches into dModel = 8, runs one transformer encoder block, then decodes the flattened token state back to a 256-pixel prefix of the original image. This is intentionally compact: it keeps the example runnable while still exercising a real patch-token transformer path. For full-image reconstruction, set reconDim := inC * inH * inW.

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 : ℕ

@[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.mkModel : API.nn.M (API.nn.Sequential σ τ)

Construct the trainable model.

The architecture lives in the public API (NN.API.Models.SelfSupervised); the example only chooses a config and trains it.

def NN.Examples.Models.Generative.Mae.loadCifarBatch (xPath yPath : System.FilePath) (nRows seed : ℕ) : IO (API.sample.Batch Float cfg.batch RealData.CifarImage RealData.CifarTarget)

Load one CIFAR minibatch as an image tensor batch.

This function deliberately stops at the data boundary: it returns CIFAR as typed image tensors. The self-supervised conversion happens in mkMaeSample, using the public SSL API, so the loader does not secretly define the model's representation.

def NN.Examples.Models.Generative.Mae.mkMaeSample (b : API.sample.Batch Float cfg.batch RealData.CifarImage RealData.CifarTarget) : API.sample.Supervised 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.trainCurve (opts : Runtime.Autograd.Torch.Options) (xPath yPath : System.FilePath) (nRows seed steps : ℕ) : IO Runtime.Training.Curve

Train and return a loss curve.

The curve is written by main using TorchLean's general training-log JSON format, so plotting and dashboard tools can consume it the same way they consume the other model examples.

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

CLI entrypoint.

Useful flags:

• --cuda runs the eager training loop on the CUDA runtime.
• --steps <n> or --epochs <n> controls optimization steps.
• --x <path> --y <path> selects custom CIFAR-style .npy arrays.
• --log <path> writes the training curve JSON.