Prediction and Training

Prediction, dataset training, checkpoint restoration, and stream training on TorchLean.Trainer.

Every method here opens a Session, drives it, and finishes it. Programs that need a different loop use trainer.open directly.

def TorchLean.Trainer.Internal.nextCyclicBatch {α : Type} (context : String) (samples : Data.SampleStream α) (cursorRef : IO.Ref ℕ) (batchSize : ℕ) : IO (Array α)

Take the next cyclic group from a finite stream, evaluating only the requested samples.

def TorchLean.Trainer.Internal.printProbes {σ τ : Shape} {trainer : Trainer σ τ} (session : trainer.Session) (probes : Array (Probe σ)) (title : String) : IO Unit

Print the predictions of every probe under a title.

def TorchLean.Trainer.Internal.loadCheckpoint {σ τ : Shape} {trainer : Trainer σ τ} (session : trainer.Session) (path? : Option System.FilePath) : IO Unit

Restore the optional checkpoint named by the training options.

def TorchLean.Trainer.Internal.saveCheckpoint {σ τ : Shape} {trainer : Trainer σ τ} (session : trainer.Session) (path? : Option System.FilePath) : IO Unit

Write the optional checkpoint named by the training options.

def TorchLean.Trainer.Internal.validateOptions (options : TrainOptions) : IO Unit

Reject training options that cannot describe a run.

def TorchLean.Trainer.Internal.runSteps {σ τ : Shape} {trainer : Trainer σ τ} (session : trainer.Session) (options : TrainOptions) (samples : Data.SampleStream (Sample.Supervised Float σ τ)) : IO Unit

Run the requested number of optimizer updates over a finite sample stream.

Batches cycle through the stream. Steps whose loss is logged read it back; every other step updates without a host readback. CUDA allocator state is sampled at the configured cadence.

def TorchLean.Trainer.predict {σ τ : Shape} (trainer : Trainer σ τ) (input : Tensor Float σ) : IO (Tensor Float τ)

Predict one input using the trainer's current model and runtime settings.

Inference before any training call. After training, use the returned trained result's trained.predict / trained.predictMany methods to predict with the trained parameters.

Example:

-- This is inference with the freshly initialized parameters, which is what makes it a useful
-- baseline to print before training starts.
def baseline (trainer : TorchLean.Trainer [2] [1]) : IO (Tensor Float [1]) :=
  trainer.predict [0.25, -0.75]

def TorchLean.Trainer.predictMany {σ τ : Shape} {batch : ℕ} (trainer : Trainer σ τ) (inputs : Tensor Float (σ.prependDim batch)) : IO (Tensor Float (τ.prependDim batch))

Predict a tensor batch using the trainer's current model and runtime settings.

def TorchLean.Trainer.train {σ τ : Shape} (trainer : Trainer σ τ) (data : Dataset σ τ) (trainOptions : TrainOptions) (probes : Array (Probe σ) := #[]) : IO (Result σ τ)

Train the model with the loss and runtime settings stored in trainer.

The signature uses Tensor Float, but the run itself executes in binary32: .native arithmetic instantiates the model over Float32 and .ieee over ExecFloat.Binary 8 23. Dataset samples are converted into that scalar as they are used, and results are read back to Float.

The result stores the trained parameters together with prediction, reporting, state access, and verification methods. trained.verify center (radius := r) (norm := .inf) checks the model that was actually trained. Add (property := .topLabel label) when the verification goal is a class label. trained.save path writes the parameters; trainer.load path data restores them.

A positive step count requires a nonempty materialized dataset.

This is trainer.open, a loop of step, and finish; the reported losses are the evaluation-mode mean losses over data before and after the loop.

Example:

-- `trained` owns the parameters the run produced; `trainer` still describes the untrained model.
def run (trainer : TorchLean.Trainer [2] [1])
    (data : Trainer.Dataset [2] [1]) : IO Unit := do
  let trained ← trainer.train data { steps := 200, logEvery := 25 }
  trained.printSummary
  let prediction ← trained.predict [0.25, -0.75]
  IO.println s!"trained(heldout) = {reprStr prediction}"

def TorchLean.Trainer.load {σ τ : Shape} (trainer : Trainer σ τ) (path : System.FilePath) (data : Dataset σ τ) : IO (Result σ τ)

Restore a checkpoint written by Result.save or Checkpoint.State.save and return it as a trained result for this trainer's model, objective, and runtime settings.

The stored Float state is cast into the run's binary32 scalar, exactly when it came from Result.save. data is evaluated once so that the result carries an honest report: steps is 0 and both losses equal the mean loss of the restored parameters on data. To continue training instead, pass loadCheckpoint? := some path in TrainOptions; the optimizer and schedule start fresh at step zero.

Example:

-- Reads back what `trained.save` wrote and re-measures the report on `data`, so both reported
-- losses are the restored model's loss and the step count is `0`.
def restore (trainer : TorchLean.Trainer [2] [1])
    (data : Trainer.Dataset [2] [1]) : IO (Trainer.Result [2] [1]) :=
  trainer.load "checkpoints/mlp.state" data

def TorchLean.Trainer.trainStream {σ τ : Shape} (trainer : Trainer σ τ) (options : Runtime.Config) (sampleAt : ℕ → Sample.Supervised Float σ τ) (evalSample : Sample.Supervised Float σ τ) (trainOptions : TrainOptions) (curveEvery : ℕ := 0) (onEval : ℕ → String → (Tensor Float σ → IO (Tensor Float τ)) → IO Unit := fun (x : ℕ) (x_1 : String) (x_2 : Tensor Float σ → IO (Tensor Float τ)) => pure ()) : IO (StreamResult σ τ)

Train a supervised model from a Float sample stream.

Generated-data examples use this when there is no fixed Dataset to hand to trainer.train. The evaluation curve is measured on evalSample; onEval receives the completed step count, a label, and a prediction function at every curve point.

Example:

-- No fixed dataset here: a sample is generated per step, and the loss curve is measured on one
-- held-out sample every `curveEvery` updates.
def sampleAt (step : Nat) : Sample.Supervised Float [1] [1] :=
  let x := step.toFloat * 0.01
  { input := [x], target := [2.0 * x] }

def run (trainer : TorchLean.Trainer [1] [1]) (runtime : Runtime.Config) :
    IO (Trainer.StreamResult [1] [1]) :=
  trainer.trainStream runtime sampleAt
    { input := [0.5], target := [1.0] }
    { steps := 500, logEvery := 100 }
    (curveEvery := 50)

def TorchLean.Trainer.trainAlternating {σ₁ τ₁ σ₂ τ₂ : Shape} (first : Trainer σ₁ τ₁) (second : Trainer σ₂ τ₂) (options : Runtime.Config) (firstSampleAt : ℕ → Sample.Supervised Float σ₁ τ₁) (secondSamplesAt : ℕ → Array (Sample.Supervised Float σ₂ τ₂)) (evalTotal : (Tensor Float σ₁ → IO (Tensor Float τ₁)) → (Tensor Float σ₂ → IO (Tensor Float τ₂)) → IO Float) (trainOptions : TrainOptions) (curveEvery : ℕ := 1) : IO (AlternatingResult σ₁ τ₁ σ₂ τ₂)

Train two supervised models from coupled Float streams.

The first model receives one supervised sample per step, while the second may receive several. evalTotal predictFirst predictSecond computes the scalar curve to record; it sees only prediction functions. Both trained results report that paired value as their before/after losses.

This alternates independent supervised updates. It does not differentiate one model's loss through the other model, so it cannot express an adversarial generator loss.