Manual Training

Dependent runners, callback composition, custom batch streams, and reporting helpers for workflows that need more control than the ordinary Trainer.new interface.

Manual Training

Direct control over runners, callbacks, data streams, and optimization steps. The ordinary Trainer interface is built from these definitions; custom training code can use them without crossing into a second API namespace.

def TorchLean.Trainer.Manual.accuracyOneHotBatched {σ : Spec.Shape} {classes batch : ℕ} {task : SeqTask (Spec.Shape.dim batch σ) (Spec.Shape.dim batch (Spec.Shape.dim classes Spec.Shape.scalar))} {α : Type} [Context α] [DecidableEq Spec.Shape] (runner : Runner α task) (samples : List (Sample.Batch α batch σ (Spec.Shape.dim classes Spec.Shape.scalar))) : IO (ℕ × ℕ)

Count correct predictions in a one-hot labeled batched dataset.

Minibatch analogue of accuracyOneHot. The task already has a leading batch axis, so the implementation scores each row independently and accumulates totals.

Returns (correct, total) where total = batch * numBatches.

structure TorchLean.Trainer.Manual.StepEvent (α : Type) : Type

Callback event fired after each training step.

• epoch : ℕ

  Current epoch number.

• step : ℕ

  Global optimizer-step counter.

• loss : α

  Loss reported for this step.

structure TorchLean.Trainer.Manual.EpochEvent : Type

Callback event fired at the end of an epoch (how many steps ran).

• epoch : ℕ

  Epoch number that just completed.

• steps : ℕ

  Number of steps executed in the epoch.

structure TorchLean.Trainer.Manual.Callbacks (α : Type) : Type

Hooks for instrumenting callback-based training loops.

Callbacks are ordinary IO hooks. They can print progress, update an in-memory curve, sample CUDA allocator state, or forward events to a project-specific metrics backend.

• onTrainStart : IO Unit

  Called once before training starts.

• onStep : StepEvent α → IO Unit

  Called after each training step.

• onEpochEnd : EpochEvent → IO Unit

  Called after each epoch.

• onTrainEnd : TrainReport α → IO Unit

  Called once after training finishes.

def TorchLean.Trainer.Manual.Callbacks.empty {α : Type} : Callbacks α

No-op callbacks.

def TorchLean.Trainer.Manual.Callbacks.append {α : Type} (a b : Callbacks α) : Callbacks α

Combine two callback collections by running them in sequence.

@[implicit_reducible]

instance TorchLean.Trainer.Manual.Callbacks.instEmptyCollection {α : Type} : EmptyCollection (Callbacks α)

∅ for callbacks: a no-op callback collection.

@[implicit_reducible]

instance TorchLean.Trainer.Manual.Callbacks.instAppend {α : Type} : Append (Callbacks α)

Callbacks form a monoid under sequential composition.

def TorchLean.Trainer.Manual.onTrainStart {α : Type} (action : IO Unit) : Callbacks α

Build callbacks that run action once at the start of training.

def TorchLean.Trainer.Manual.onStep {α : Type} (f : StepEvent α → IO Unit) : Callbacks α

Build callbacks that observe every training step.

def TorchLean.Trainer.Manual.cudaMemWatchCallbacks {α : Type} (opts : Runtime.Autograd.Torch.Options) (watchEvery totalSteps : ℕ) : IO (Callbacks α)

Build a training callback that samples the CUDA allocator at a fixed step cadence.

The callback owns a small IO.Ref for the previous sample, so examples can compose it with ordinary loss-logging callbacks without threading allocator state through their training loops.

def TorchLean.Trainer.Manual.onEpochEnd {α : Type} (f : EpochEvent → IO Unit) : Callbacks α

Build callbacks that run at the end of each epoch.

def TorchLean.Trainer.Manual.onTrainEnd {α : Type} (f : TrainReport α → IO Unit) : Callbacks α

Build callbacks that run once at the end of training, with the final report.

structure TorchLean.Trainer.Manual.StepBatchStream (α : Type) (inputShapes : List Spec.Shape) : Type

Step-indexed source of already-collated module inputs.

TorchLean.Data.batchLoader is the right interface when the data is a finite supervised dataset. Other training jobs draw batches from a rule or an external source: replay buffers, collocation samplers, synthetic scale inputs, or file-backed sequence windows. StepBatchStream is the direct stream interface for those cases.

The stream is still fully typed: each produced sample is a _root_.Runtime.Autograd.Torch.TList matching the module's inputShapes. The training loop below is model-agnostic and only assumes that the module can run forward and stepWith on those samples.

• sample : ℕ → IO (Runtime.Autograd.Torch.TList α inputShapes)

  Produce the input sample used at logical optimizer step step.

def TorchLean.Trainer.Manual.StepBatchStream.fixed {α : Type} {inputShapes : List Spec.Shape} (x : Runtime.Autograd.Torch.TList α inputShapes) : StepBatchStream α inputShapes

Constant stream for fixed-batch overfit runs and fixed-sample training jobs.

def TorchLean.Trainer.Manual.StepBatchStream.ofFn {α : Type} {inputShapes : List Spec.Shape} (f : ℕ → Runtime.Autograd.Torch.TList α inputShapes) : StepBatchStream α inputShapes

Build a stream from a pure step-indexed sample function.

def TorchLean.Trainer.Manual.StepBatchStream.cycle {α : Type} {inputShapes : List Spec.Shape} (xs : List (Runtime.Autograd.Torch.TList α inputShapes)) (h : xs ≠ []) : StepBatchStream α inputShapes

Cycle through a nonempty list of samples.

This adapter lets list-backed datasets use the step-stream trainer. The explicit nonempty proof keeps empty datasets from turning into silent modulo-by-zero behavior.

def TorchLean.Trainer.Manual.withMode {σ τ : Spec.Shape} {task : SeqTask σ τ} {α : Type} [Context α] [DecidableEq Spec.Shape] {β : Type} (runner : Runner α task) (value : Runtime.Autograd.TorchLean.NN.Mode) (action : IO β) : IO β

Run an action with the runner temporarily switched to value mode.

Use this for callback-based validation passes during training.

def TorchLean.Trainer.Manual.meanLossModuleLoader {σ τ : Spec.Shape} {n : ℕ} {paramShapes : List Spec.Shape} {α : Type} [Context α] [DecidableEq Spec.Shape] [ToString α] (module : Runtime.Autograd.TorchLean.ScalarModule α paramShapes [Spec.Shape.dim n σ, Spec.Shape.dim n τ]) (loader : Data.BatchLoader α n σ τ) : IO α

Mean loss for an already-instantiated scalar module over a typed minibatch loader.

General streaming evaluation path used by the runtime examples. It is not CIFAR-specific: any supervised task whose loss module consumes [dim n σ, dim n τ] can use the same loader. The loader stores ordinary per-example samples (x : σ, y : τ); this definition asks TorchLean.Data.epoch for raw minibatches and calls TorchLean.Data.collateSupervised to build one shape-typed batch at a time.

Two details matter for larger examples:

• We force shuffle := false for evaluation so before/after metrics are deterministic.
• We do not call TorchLean.Data.BatchLoader.batchDataset, because that would materialize every collated minibatch at once. Streaming keeps the same API usable for image, sequence, and scientific ML examples where the batch tensors are much larger than small tabular datasets.

def TorchLean.Trainer.Manual.meanLossBatchLoader {σ τ : Spec.Shape} {n : ℕ} {task : SeqTask (Spec.Shape.dim n σ) (Spec.Shape.dim n τ)} {α : Type} [Context α] [DecidableEq Spec.Shape] [ToString α] (runner : Runner α task) (loader : Data.BatchLoader α n σ τ) : IO α

Mean loss over a typed minibatch loader through a Trainer.Manual.Runner.

Runner-facing form of meanLossModuleLoader. Use it when the example is built around Trainer.Manual.run, task modes, and the proof layer trainer abstraction. Use meanLossModuleLoader directly when the example has already instantiated a runtime TorchLean.Module.ScalarModule, which is the common fast path for CUDA examples.

def TorchLean.Trainer.Manual.accuracyOneHotBatchLoader {σ : Spec.Shape} {classes batch : ℕ} {task : SeqTask (Spec.Shape.dim batch σ) (Spec.Shape.dim batch (Spec.Shape.dim classes Spec.Shape.scalar))} {α : Type} [Context α] [DecidableEq Spec.Shape] (runner : Runner α task) (loader : Data.BatchLoader α batch σ (Spec.Shape.dim classes Spec.Shape.scalar)) : IO (ℕ × ℕ)

One-hot accuracy over a typed minibatch loader without materializing all collated batches.

def TorchLean.Trainer.Manual.trainModuleLoaderWith {σ τ : Spec.Shape} {n : ℕ} {paramShapes : List Spec.Shape} {α : Type} [Context α] [DecidableEq Spec.Shape] [ToString α] (module : Runtime.Autograd.TorchLean.ScalarModule α paramShapes [Spec.Shape.dim n σ, Spec.Shape.dim n τ]) (optimizer : Runtime.Autograd.TorchLean.Optimizer α paramShapes) (epochs : ℕ) (loader : Data.BatchLoader α n σ τ) (callbacks : Callbacks α := Callbacks.empty) : IO (TrainReport α × Data.BatchLoader α n σ τ)

Train a runtime scalar module from a typed minibatch loader.

Shared real epoch loop for model examples that already have a runtime module, including CUDA runs. It mirrors the PyTorch structure:

1. create an optimizer state for the module parameters;
2. for each epoch, ask the general TorchLean.Data.batchLoader for shuffled raw batches;
3. collate each raw batch into a shape-typed (xBatch, yBatch) sample;
4. report the scalar loss through callbacks;
5. run forward/backward/optimizer.step through TorchLean.Module.stepWith.

The function is polymorphic in the input shape σ, target shape τ, batch size n, scalar type α, parameter shapes, and optimizer. It is not image-specific. CNN, ResNet, ViT, MLP, sequence, operator-learning, and future model examples should all be able to use this path whenever their supervised loss module has input shapes [dim n σ, dim n τ].

def TorchLean.Trainer.Manual.trainModuleLoaderStepsWith {σ τ : Spec.Shape} {n : ℕ} {paramShapes : List Spec.Shape} {α : Type} [Context α] [DecidableEq Spec.Shape] [ToString α] (module : Runtime.Autograd.TorchLean.ScalarModule α paramShapes [Spec.Shape.dim n σ, Spec.Shape.dim n τ]) (optimizer : Runtime.Autograd.TorchLean.Optimizer α paramShapes) (steps : ℕ) (loader : Data.BatchLoader α n σ τ) (callbacks : Callbacks α := Callbacks.empty) : IO (TrainReport α × Data.BatchLoader α n σ τ)

Train a runtime scalar module for exactly steps optimizer updates.

trainModuleLoaderWith above is epoch-based: each unit means one full pass over the loader. This variant is update-based, which is the convention used by runnable examples that expose a --steps flag.

The loop still draws shuffled minibatches from TorchLean.Data.batchLoader epoch by epoch, but it stops as soon as the requested number of optimizer updates has run. The returned loader is the next loader state, so callers can continue training from the next shuffled epoch if they want to.

def TorchLean.Trainer.Manual.trainModuleStreamStepsWith {inputShapes paramShapes : List Spec.Shape} {α : Type} [Context α] [DecidableEq Spec.Shape] [ToString α] (module : Runtime.Autograd.TorchLean.ScalarModule α paramShapes inputShapes) (optimizer : Runtime.Autograd.TorchLean.Optimizer α paramShapes) (steps : ℕ) (stream : StepBatchStream α inputShapes) (callbacks : Callbacks α := Callbacks.empty) : IO (TrainReport α)

Train a scalar module from a step-indexed batch stream.

Shared loop for workloads whose batches are produced step by step rather than by one finite TorchLean.Data.batchLoader epoch:

• RL algorithms can sample replay or rollout batches,
• PDE examples can resample collocation points,
• generated workloads can stream synthetic inputs without storing a dataset.

The function is generic in inputShapes. It does not know whether the sample is [x, y], [state, action, target], or []; it only asks the stream for the next typed input list and then runs the same forward/backward/optimizer.step machinery as the loader-based trainer.

def TorchLean.Trainer.Manual.trainModuleStreamStepsReport {inputShapes paramShapes : List Spec.Shape} {α : Type} [Context α] [DecidableEq Spec.Shape] [ToString α] (module : Runtime.Autograd.TorchLean.ScalarModule α paramShapes inputShapes) (optimizer : Runtime.Autograd.TorchLean.Optimizer α paramShapes) (opts : Runtime.Autograd.Torch.Options) (steps : ℕ) (stream : StepBatchStream α inputShapes) (cudaMemWatch : ℕ := 0) (extraCallbacks : Callbacks α := Callbacks.empty) : IO (TrainReport α)

Report-oriented stream-training entrypoint.

Callers pass the module, optimizer, runtime options, step count, and stream, and get standard before/after reporting plus CUDA memory watching.

def TorchLean.Trainer.Manual.trainModuleStreamStepsCurveFloat {inputShapes paramShapes : List Spec.Shape} (module : Runtime.Autograd.TorchLean.ScalarModule Float paramShapes inputShapes) (optimizer : Runtime.Autograd.TorchLean.Optimizer Float paramShapes) (opts : Runtime.Autograd.Torch.Options) (steps : ℕ) (stream : StepBatchStream Float inputShapes) (cudaMemWatch : ℕ := 0) (extraCallbacks : Callbacks Float := Callbacks.empty) : IO (TrainReport Float × Training.Curve)

Float stream trainer that records a per-step loss curve.

Generated and file-backed batches do not always have one finite loader to summarize. This entrypoint keeps their training curves in the same JSON format as the supervised examples.

def TorchLean.Trainer.Manual.trainLoaderWith {σ τ : Spec.Shape} {n : ℕ} {task : SeqTask (Spec.Shape.dim n σ) (Spec.Shape.dim n τ)} {α : Type} [Context α] [DecidableEq Spec.Shape] [ToString α] [Runtime.FromFloat α] (runner : Runner α task) (cfg : LoaderTrainConfig) (loader : Data.BatchLoader α n σ τ) (callbacks : Callbacks α := Callbacks.empty) : IO (TrainReport α × Data.BatchLoader α n σ τ)

Train from a runner-backed loader with explicit callbacks instead of inline printing in example code.

Runner-facing public path for PyTorch-style custom loops:

• keep the optimizer/scheduler logic in the library,
• inject logging, evaluation, and prediction reporting through callbacks.

This path keeps the Runner abstraction, including task modes and scheduler support. For CUDA-heavy entrypoints that already have a TorchLean.Module.ScalarModule, prefer trainModuleLoaderWith; both paths consume the same general TorchLean.Data.batchLoader.

Small Reporting Helpers (IO)

These definitions factor out common "print a loss/accuracy table" patterns for runnable model commands. They do not affect semantics: they only call the underlying runner functions and print human-facing summaries. Public examples should reach them through Trainer.Manual only when the ordinary Trainer.new / trainer.train API is too small for the example.

def TorchLean.Trainer.Manual.Report.reportProbes {β : Type} (title : String) (probes : List β) (lineOf : β → IO String) : IO Unit

Print a titled list of named report lines.

def TorchLean.Trainer.Manual.Report.reportMeanLoss {σ τ : Spec.Shape} {task : SeqTask σ τ} {α : Type} [Context α] [DecidableEq Spec.Shape] [ToString α] (runner : Runner α task) (dataset : Runtime.Autograd.Train.Dataset (Sample.Supervised α σ τ)) (label : String) : IO Unit

Convenience: mean loss on a dataset, printed with a label.

def TorchLean.Trainer.Manual.Report.reportMeanLossLoader {σ τ : Spec.Shape} {batch : ℕ} {task : SeqTask (Spec.Shape.dim batch σ) (Spec.Shape.dim batch τ)} {α : Type} [Context α] [DecidableEq Spec.Shape] [ToString α] (runner : Runner α task) (loader : Data.BatchLoader α batch σ τ) (label : String) : IO Unit

Convenience: mean loss on a typed minibatch loader, streamed batch by batch.

def TorchLean.Trainer.Manual.Report.reportMeanLossModuleLoader {σ τ : Spec.Shape} {batch : ℕ} {paramShapes : List Spec.Shape} {α : Type} [Context α] [DecidableEq Spec.Shape] [ToString α] (module : Runtime.Autograd.TorchLean.ScalarModule α paramShapes [Spec.Shape.dim batch σ, Spec.Shape.dim batch τ]) (loader : Data.BatchLoader α batch σ τ) (label : String) : IO Unit

Convenience: mean loss on a typed minibatch loader for an already-instantiated runtime module.

Use this in direct CUDA/runtime examples to avoid building a Runner only for logging. The data path is still the same public loader path: TorchLean.Data.batchLoader plus TorchLean.Data.collateSupervised.

def TorchLean.Trainer.Manual.Report.reportClassProbes {σ : Spec.Shape} {classes : ℕ} {task : SeqTask σ (Spec.Shape.dim classes Spec.Shape.scalar)} {α : Type} [Context α] [DecidableEq Spec.Shape] [ToString α] (runner : Runner α task) (probes : List (String × Spec.Tensor α σ × ℕ)) (title : String := "predictions") (includeLogits : Bool := false) : IO Unit

Report predicted classes on a list of named inputs.

Each entry is (name, x, expectedClass). If includeLogits := true, also prints the raw model outputs.

def TorchLean.Trainer.Manual.Report.reportClassProbesBatchedFromSingle {σ : Spec.Shape} {classes batch : ℕ} {task : SeqTask (Spec.Shape.dim batch σ) (Spec.Shape.dim batch (Spec.Shape.dim classes Spec.Shape.scalar))} {α : Type} [Context α] [DecidableEq Spec.Shape] [ToString α] (runner : Runner α task) (probes : List (String × Spec.Tensor α σ × ℕ)) (title : String := "predictions") (includeLogits : Bool := false) : IO Unit

Report predicted classes on a list of named inputs, for a batched model.

This expects inputs of the unbatched input shape σ and replicates each one across the batch axis, then reports the prediction for row 0.

def TorchLean.Trainer.Manual.Report.reportLossAccuracyOneHot {σ : Spec.Shape} {classes : ℕ} {task : SeqTask σ (Spec.Shape.dim classes Spec.Shape.scalar)} {α : Type} [Context α] [DecidableEq Spec.Shape] [ToString α] (runner : Runner α task) (dataset : Runtime.Autograd.Train.Dataset (Sample.Supervised α σ (Spec.Shape.dim classes Spec.Shape.scalar))) (label : String) : IO Unit

Convenience: mean loss + one-hot accuracy on a dataset, printed with a label.

def TorchLean.Trainer.Manual.Report.reportLossAccuracyOneHotBatched {σ : Spec.Shape} {classes batch : ℕ} {task : SeqTask (Spec.Shape.dim batch σ) (Spec.Shape.dim batch (Spec.Shape.dim classes Spec.Shape.scalar))} {α : Type} [Context α] [DecidableEq Spec.Shape] [ToString α] (runner : Runner α task) (dataset : Runtime.Autograd.Train.Dataset (Sample.Batch α batch σ (Spec.Shape.dim classes Spec.Shape.scalar))) (label : String) : IO Unit

Batched variant of reportLossAccuracyOneHot.

def TorchLean.Trainer.Manual.Report.reportLossAccuracyOneHotLoader {σ : Spec.Shape} {classes batch : ℕ} {task : SeqTask (Spec.Shape.dim batch σ) (Spec.Shape.dim batch (Spec.Shape.dim classes Spec.Shape.scalar))} {α : Type} [Context α] [DecidableEq Spec.Shape] [ToString α] (runner : Runner α task) (loader : Data.BatchLoader α batch σ (Spec.Shape.dim classes Spec.Shape.scalar)) (label : String) : IO Unit

Loader variant of reportLossAccuracyOneHotBatched, streaming through minibatches.

def TorchLean.Trainer.Manual.trainModuleLoaderStepsReport {σ τ : Spec.Shape} {n : ℕ} {paramShapes : List Spec.Shape} {α : Type} [Context α] [DecidableEq Spec.Shape] [ToString α] (module : Runtime.Autograd.TorchLean.ScalarModule α paramShapes [Spec.Shape.dim n σ, Spec.Shape.dim n τ]) (optimizer : Runtime.Autograd.TorchLean.Optimizer α paramShapes) (opts : Runtime.Autograd.Torch.Options) (steps : ℕ) (loader : Data.BatchLoader α n σ τ) (cudaMemWatch : ℕ := 0) (extraCallbacks : Callbacks α := Callbacks.empty) : IO (TrainReport α × Data.BatchLoader α n σ τ)

Train a runtime module for a fixed number of optimizer updates with the standard runtime reports.

Common path for direct-module training, not example-only code. It composes the generic step loop with before/after mean-loss reporting and CUDA allocator telemetry, while still accepting extra callbacks for projects that want their own metrics, validation, or tracing.

def TorchLean.Trainer.Manual.trainModuleLoaderStepsCurveFloat {σ τ : Spec.Shape} {n : ℕ} {paramShapes : List Spec.Shape} (module : Runtime.Autograd.TorchLean.ScalarModule Float paramShapes [Spec.Shape.dim n σ, Spec.Shape.dim n τ]) (optimizer : Runtime.Autograd.TorchLean.Optimizer Float paramShapes) (opts : Runtime.Autograd.Torch.Options) (steps : ℕ) (loader : Data.BatchLoader Float n σ τ) (cudaMemWatch : ℕ := 0) (extraCallbacks : Callbacks Float := Callbacks.empty) : IO (TrainReport Float × Data.BatchLoader Float n σ τ × Training.Curve)

Float-specialized module training that also records a scalar loss curve.

The training loop itself is the same as trainModuleLoaderStepsReport; this entrypoint adds the standard Curve callback used by JSON logs and website widgets.

def TorchLean.Trainer.Manual.trainModuleLoaderStepsLoggedFloat {σ τ : Spec.Shape} {n : ℕ} {paramShapes : List Spec.Shape} (module : Runtime.Autograd.TorchLean.ScalarModule Float paramShapes [Spec.Shape.dim n σ, Spec.Shape.dim n τ]) (optimizer : Runtime.Autograd.TorchLean.Optimizer Float paramShapes) (opts : Runtime.Autograd.Torch.Options) (steps : ℕ) (loader : Data.BatchLoader Float n σ τ) (log : Training.LogDestination) (title : String) (notes : Array String := #[]) (seriesName : String := "loss") (cudaMemWatch : ℕ := 0) : IO (TrainReport Float × Data.BatchLoader Float n σ τ)

Train a Float runtime module, write a standard scalar-curve log, and return the train report.

High-level path used by runnable training commands. The caller provides the model, optimizer, loader, runtime options, and metadata notes; the library owns the callback composition, CUDA telemetry, before/after reports, and JSON curve emission.