Training

This page covers loss functions, optimizers, and the training loop. For ownership management and scoped { }, see Core Concepts.

Loss Functions

Each loss function takes two tensors and returns a Tensor. Common choices include Loss.crossEntropy for classification and Loss.mse for regression:

let loss = Loss.crossEntropy predictions targets
printfn "Loss: %.4f" (loss.item ())

See the Loss API reference for the full list (mse, nll, crossEntropy, binaryCrossEntropyWithLogit, l1, smoothL1, klDiv).

Optimizers

SGD

Stochastic Gradient Descent:

let vars = Model.trainableVars model
let opt = SGD.create 0.01 vars

AdamW

Adam with weight decay:

let vars = Model.trainableVars model
let opt = AdamW.createWithLr 0.001 vars

Use AdamW.create for full control over hyperparameters:

let adamw = AdamW.create {
    ParamsAdamW.defaultParams with
        Lr = 0.001
        Beta1 = 0.9
        Beta2 = 0.999
} vars

Optimizer Members

Both SGD and AdamW expose these member methods directly:

• step() -- update parameters using current gradients.
• zeroGrad() -- clear accumulated gradients. Call before each training step (or skip for gradient accumulation).
• learningRate() -- return the current learning rate.
• setLearningRate lr -- set a new learning rate.
• toOps() -- create an OptimizerOps record for use with Checkpoint.save / Checkpoint.load.

The Training Loop

A typical training loop uses scoped { } to free intermediate tensors after each batch. Without scoped { }, tensors created during the forward and backward passes accumulate until the garbage collector runs, which can cause high native memory usage:

for epoch in 1..epochs do
    for batch in trainLoader do
        scoped {
            let x, y = batch
            opt.zeroGrad ()
            let pred = model.forward x
            let loss = Loss.crossEntropy pred y
            loss.backward ()
            opt.step ()
        }

    if epoch % 10 = 0 then
        printfn "epoch %d" epoch

See Core Concepts -- Memory Management for details on scoped { }.

Evaluation

Wrap inference in Toro.noGrad to disable gradient tracking. This reduces memory use and speeds up the forward pass:

Toro.noGrad (fun () ->
    let pred = model.forward testX
    let loss = Loss.crossEntropy pred testY
    printfn "Test loss: %.4f" (loss.item ())

    let argmax = pred.argmax 1
    let eq = (argmax .=. testY).toDType F32
    let total = eq.sumAll ()
    printfn "Accuracy: %.2f%%" (total.item () / float testY.ElemCount * 100.0)
)

Learning Rate Scheduling

LrSchedule is a discriminated union with cases StepDecay, Exponential, CosineAnnealing, and LinearWarmup. Create a Scheduler by pairing a schedule with a setLr callback. See the Scheduler API reference for all schedule types.

let sched = Scheduler.create (CosineAnnealing(100, 0.0)) opt.setLearningRate (opt.learningRate ())

for epoch in 1..epochs do
    opt.zeroGrad ()
    let pred = model.forward x
    let loss = Loss.mse pred y
    loss.backward ()
    opt.step ()
    Scheduler.step sched

You can also use setLearningRate directly for custom schedules:

if epoch % 100 = 0 then
    opt.setLearningRate (opt.learningRate () * 0.5)

Gradient Clipping

Clip gradients to prevent exploding gradients during training:

let vars = Model.trainableVars model

// Clip by global L2 norm (returns the total norm before clipping)
let totalNorm = Clip.gradNorm 1.0 vars

// Or clip each gradient element to [-0.5, 0.5]
Clip.gradValue 0.5 vars

Call gradient clipping after loss.backward() and before opt.step():

opt.zeroGrad ()
let pred = model.forward x
let loss = Loss.crossEntropy pred y
loss.backward ()
let _ = Clip.gradNorm 1.0 vars
opt.step ()

Metrics

The Metrics module provides classification metrics (accuracy, accuracyFromLogits, precision, recall, f1). See the Metrics API reference for details.

Toro.noGrad (fun () ->
    let pred = model.forward testX
    let acc = Metrics.accuracyFromLogits pred testY
    printfn "Accuracy: %.2f%%" (acc * 100.0)
)

Checkpointing

Full checkpoint (model + optimizer + epoch)

Checkpoint.save and Checkpoint.load persist the complete training state: model parameters, optimizer internal state (e.g. AdamW momentum), and the epoch number.

for epoch in 1..epochs do
    opt.zeroGrad ()
    let pred = model.forward x
    let loss = Loss.crossEntropy pred y
    loss.backward ()
    opt.step ()

    if epoch % 50 = 0 then
        Checkpoint.save model (opt.toOps ()) epoch $"checkpoints/epoch{epoch}"

Resume training from a checkpoint:

let resumedEpoch = Checkpoint.load model (opt.toOps ()) "checkpoints/epoch200"
printfn "Resumed from epoch %d" resumedEpoch

for epoch in (resumedEpoch + 1)..epochs do
    // training continues with restored optimizer state

Checkpoint.load also restores the learning rate stored at save time.

Model-only save/load

For saving model weights without optimizer state, use Model.save and Model.loadInto:

Model.save model "weights/best.safetensors"
let report = Model.loadInto model "weights/best.safetensors" Strict

Model.loadInto matches parameters by name, validates shape and dtype, copies values into the model tensors, and returns a LoadReport. Only the tensors required by the model are loaded into memory. Use Strict to error on missing keys, unexpected keys, or shape/dtype mismatches, or Lenient to allow partial loading.

Pre-trained Models

See Hub for downloading and loading pre-trained weights from Hugging Face Hub.