TorchLean API

NN.API.Trainer.Session

Training Sessions #

trainer.train owns the optimizer loop. A Session is the same trainer with the loop handed to the caller:

let session ← trainer.open
for step in [0:steps] do
  let loss ← session.step (sampleAt step)
  if step % 100 = 0 then IO.println s!"step {step}: loss={loss}"
let after ← session.eval data
let trained ← session.finish { before, after }

Opening a session instantiates the model under the trainer's runtime settings and binds the trainer's optimizer. step and stepBatch apply one update and return the loss; update and updateBatch apply the same update without reading the loss; predict, loss, and eval run the current parameters in evaluation mode; state, save, and load read or replace the parameters; finish snapshots the current state as an ordinary Result.

Every public signature uses Float. The run executes in the binary32 scalar selected by RunConfig.arithmetic, and values are converted at the boundary. Trainer.train is implemented as open, a loop, and finish.

structure TorchLean.Trainer.Session {σ τ : Shape} (trainer : Trainer σ τ) :

A trainer with its model instantiated and the optimizer loop owned by the caller.

Obtain one with trainer.open. The session type is indexed by the trainer it runs, so the state layout nn.stateShapes trainer.model is available to state and load. Mode handling is implicit: updates run stateful layers in training mode, while predict, loss, and eval run them in evaluation mode.

Instances For
    opaque TorchLean.Trainer.Session.Internal.create {σ τ : Shape} (trainer : Trainer σ τ) (stepBatch : Array (Sample.Supervised Float σ τ)IO Float) (update : Array (Sample.Supervised Float σ τ)IO Unit) (steps : IO ) (loss : Sample.Supervised Float σ τIO Float) (predict : Tensor Float σIO (Tensor Float τ)) (state : IO (nn.State Float (Runtime.Autograd.Model.Layers.Seq.stateShapes trainer.model))) (setState : nn.State Float (Runtime.Autograd.Model.Layers.Seq.stateShapes trainer.model)IO Unit) (finish : Training.LossProgress FloatIO (Result σ τ)) :
    trainer.Session

    Construct a session at the trainer implementation boundary.

    Replace the live parameters and buffers from host Float tensors.

    def TorchLean.Trainer.Session.trainer {σ τ : Shape} {trainer : Trainer σ τ} :
    trainer.SessionTrainer σ τ

    The trainer whose model, objective, and runtime settings this session runs.

    Instances For
      opaque TorchLean.Trainer.Session.step {σ τ : Shape} {trainer : Trainer σ τ} (session : trainer.Session) (sample : Sample.Supervised Float σ τ) :

      Apply one optimizer update on a single sample and return its loss.

      Example:

      -- One update per call, returning that sample's loss. Batching several samples into a single
      -- update is `stepBatch`, not repeated `step` calls.
      def descend (trainer : TorchLean.Trainer [2] [1]) : IO (Trainer.Result [2] [1]) := do
        let session ← trainer.open
        let sample : Sample.Supervised Float [2] [1] := { input := [1.0, 0.0], target := [1.0] }
        let before ← session.loss sample
        for _ in List.range 100 do
          let _ ← session.step sample
        let after ← session.loss sample
        session.finish { before := before, after := after }
      
      opaque TorchLean.Trainer.Session.stepBatch {σ τ : Shape} {trainer : Trainer σ τ} (session : trainer.Session) (batch : Array (Sample.Supervised Float σ τ)) :

      Apply one optimizer update on a nonempty batch and return the mean loss.

      The per-sample gradients are averaged at the same parameter point, so the batch counts as a single optimizer step. An empty batch is rejected.

      opaque TorchLean.Trainer.Session.update {σ τ : Shape} {trainer : Trainer σ τ} (session : trainer.Session) (sample : Sample.Supervised Float σ τ) :

      Apply one optimizer update without returning the loss.

      Use this for unlogged steps in custom loops. Native CUDA optimizers keep gradients and moments on device and avoid the loss readback performed by step.

      opaque TorchLean.Trainer.Session.updateBatch {σ τ : Shape} {trainer : Trainer σ τ} (session : trainer.Session) (batch : Array (Sample.Supervised Float σ τ)) :

      Apply one mean-gradient update to a nonempty batch without returning the loss.

      This shares the optimizer history and step counter used by step, stepBatch, and update.

      opaque TorchLean.Trainer.Session.steps {σ τ : Shape} {trainer : Trainer σ τ} (session : trainer.Session) :

      Number of optimizer updates applied so far.

      opaque TorchLean.Trainer.Session.loss {σ τ : Shape} {trainer : Trainer σ τ} (session : trainer.Session) (sample : Sample.Supervised Float σ τ) :

      Loss of one sample under the current parameters in evaluation mode.

      opaque TorchLean.Trainer.Session.predict {σ τ : Shape} {trainer : Trainer σ τ} (session : trainer.Session) (input : Tensor Float σ) :

      Run one input through the current parameters in evaluation mode.

      Read the current parameters and buffers as host Float tensors.

      opaque TorchLean.Trainer.Session.finish {σ τ : Shape} {trainer : Trainer σ τ} (session : trainer.Session) (loss : Training.LossProgress Float) :
      IO (Result σ τ)

      Snapshot the current parameters as a trained Result.

      The report records the number of updates applied so far, the runtime arithmetic, and the losses supplied by the caller, typically from eval before and after the loop. Later session updates or checkpoint loads do not change the result's parameters, predictions, or verification.

      Example:

      -- Keep a model snapshot with the losses measured for it. The session can continue training
      -- without changing this result.
      def package {trainer : TorchLean.Trainer [2] [1]}
          (session : Trainer.Session trainer) (before after : Float) :
          IO (Trainer.Result [2] [1]) :=
        session.finish { before := before, after := after }
      
      def TorchLean.Trainer.Session.predictMany {σ τ : Shape} {batch : } {trainer : Trainer σ τ} (session : trainer.Session) (inputs : Tensor Float (σ.prependDim batch)) :
      IO (Tensor Float (τ.prependDim batch))

      Run several inputs through the current parameters in evaluation mode.

      Instances For
        def TorchLean.Trainer.Session.meanLoss {σ τ : Shape} {trainer : Trainer σ τ} (session : trainer.Session) (samples : Data.SampleStream (Sample.Supervised Float σ τ)) :

        Mean loss over a finite sample stream in evaluation mode; 0 for an empty stream.

        Instances For
          def TorchLean.Trainer.Session.eval {σ τ : Shape} {trainer : Trainer σ τ} (session : trainer.Session) (data : Dataset σ τ) :

          Mean loss over a dataset in evaluation mode.

          Instances For
            def TorchLean.Trainer.Session.save {σ τ : Shape} {trainer : Trainer σ τ} (session : trainer.Session) (path : System.FilePath) :

            Write the current parameters with Checkpoint.State.save.

            Instances For
              def TorchLean.Trainer.Session.load {σ τ : Shape} {trainer : Trainer σ τ} (session : trainer.Session) (path : System.FilePath) :

              Replace the current model state from a checkpoint written by save or Result.save.

              Optimizer history and the completed-step count are retained. Open a new session before loading when training should restart with a fresh optimizer and schedule.

              Instances For

                Build a session for one runtime scalar.

                The scalar must support host readback and the verifier's bound arithmetic so that the finished result can predict, save, and verify.

                open calls this at two concrete scalars, and without nospecialize the compiler generates a specialized copy of the whole runner and stepper stack for each of them. That accounted for 14 of the 22 seconds this module used to take. Session setup runs once per training run, so the generic version costs nothing that matters.

                Instances For
                  def TorchLean.Trainer.open {σ τ : Shape} (trainer : Trainer σ τ) (scheduler : Option Scheduler.Config := none) :
                  IO trainer.Session

                  Instantiate the model under the trainer's runtime settings and hand the optimizer loop to the caller.

                  The trainer's optimizer is bound immediately; scheduler optionally adjusts its learning rate by completed step. Optimizer and scheduler settings must remain valid after binary32 conversion and are checked before the model is instantiated. CUDA execution requires .native arithmetic, and .complex arithmetic is not supported by supervised training.

                  Example:

                  -- A session holds the instantiated model and the bound optimizer. `trainer.train` is exactly this
                  -- call, a loop of `step`, and a `finish`, so opening a session is how you take that loop over.
                  def stepOnce (trainer : TorchLean.Trainer [2] [1]) : IO Float := do
                    let session ← trainer.open
                    session.step { input := [1.0, 0.0], target := [1.0] }
                  
                  Instances For