Training

Training loops in AodhML are explicit but ergonomic. The language provides primitives for forward passes, backward passes, and parameter updates.

Prototype Status: Training primitives are planned for v0.4.0 (autodiff). The examples below show the intended API.

Training Loop

import "std:nn"
import "std:optim"
import "std:loss"

fn train(
    model: Model,
    data: Dataset<Tensor<f32>, Tensor<f32>>,
    epochs: u32,
    lr: f32,
) {
    let optimizer = Adam::new(lr: lr)

    for epoch in 0..epochs {
        mut total_loss = 0.0

        for batch in data.batches(32) {
            let (x, y) = batch

            // Forward
            let pred = model.forward(x)
            let loss = cross_entropy(pred, y)

            // Backward
            let grads = loss.backward()

            // Update
            optimizer.step(model.parameters(), grads)

            total_loss = total_loss + loss.value()
        }

        let avg_loss = total_loss / data.batches(32).len() as f32
        println("Epoch ", epoch, " — Loss: ", avg_loss)
    }
}

Optimizers

OptimizerStatus
SGDPlanned v0.4.0
AdamPlanned v0.4.0
AdamWPlanned v0.5.0
RMSpropPlanned v0.5.0

Loss Functions

LossStatus
mse_lossPlanned v0.4.0
cross_entropyPlanned v0.4.0
nll_lossPlanned v0.5.0