Trainer API

The Trainer struct provides a high-level abstraction for training neural networks with full callback support, automatic metrics tracking, and gradient management.

Overview

#![allow(unused)]
fn main() {
use entrenar::train::{Trainer, TrainConfig, Batch, MSELoss, EarlyStopping};
use entrenar::optim::Adam;
use entrenar::Tensor;

// Create trainer
let params = vec![Tensor::zeros(784 * 128, true)];
let optimizer = Adam::new(0.001, 0.9, 0.999, 1e-8);
let config = TrainConfig::default();

let mut trainer = Trainer::new(params, Box::new(optimizer), config);
trainer.set_loss(Box::new(MSELoss));

// Add callbacks
trainer.add_callback(EarlyStopping::new(5, 0.001));

// Train
let result = trainer.train(100, || batches.clone(), |x| model.forward(x));
}

Trainer Struct

#![allow(unused)]
fn main() {
pub struct Trainer {
    params: Vec<Tensor>,           // Model parameters
    optimizer: Box<dyn Optimizer>, // Optimizer instance
    loss_fn: Option<Box<dyn LossFn>>, // Loss function
    config: TrainConfig,           // Training configuration
    pub metrics: MetricsTracker,   // Metrics tracking
    callbacks: CallbackManager,    // Callback system
    best_loss: Option<f32>,        // Best loss achieved
    start_time: Option<Instant>,   // Training start time
}
}

Creating a Trainer

#![allow(unused)]
fn main() {
let trainer = Trainer::new(params, optimizer, config);
}

Parameters:

• params: Vec<Tensor> - Model parameters to optimize (must have requires_grad = true)
• optimizer: Box<dyn Optimizer> - Optimizer instance (SGD, Adam, AdamW)
• config: TrainConfig - Training configuration

Setting the Loss Function

#![allow(unused)]
fn main() {
trainer.set_loss(Box::new(MSELoss));
// or
trainer.set_loss(Box::new(CrossEntropyLoss));
}

The loss function must be set before calling train() or train_step().

Adding Callbacks

#![allow(unused)]
fn main() {
use entrenar::train::{EarlyStopping, CheckpointCallback, ProgressCallback, MonitorCallback};

trainer.add_callback(EarlyStopping::new(5, 0.001));
trainer.add_callback(CheckpointCallback::new("./checkpoints"));
trainer.add_callback(ProgressCallback::new(10));
trainer.add_callback(MonitorCallback::new());
}

See Callback System for details on available callbacks.

Training Methods

train() - Full Training Loop

The primary method for training with full callback support:

#![allow(unused)]
fn main() {
pub fn train<F, B, I>(
    &mut self,
    max_epochs: usize,
    batch_fn: B,
    forward_fn: F,
) -> TrainResult
where
    F: Fn(&Tensor) -> Tensor,
    B: Fn() -> I,
    I: IntoIterator<Item = Batch>,
}

Parameters:

• max_epochs - Maximum number of epochs to train
• batch_fn - Function that returns batches for each epoch
• forward_fn - Model forward pass (inputs → predictions)

Returns: TrainResult with training outcome

Example:

#![allow(unused)]
fn main() {
let batches = vec![
    Batch::new(inputs1, targets1),
    Batch::new(inputs2, targets2),
];

let result = trainer.train(
    100,                          // max epochs
    || batches.clone(),           // batch function
    |x| model.forward(x),         // forward function
);

println!("Final epoch: {}", result.final_epoch);
println!("Final loss: {:.4}", result.final_loss);
println!("Best loss: {:.4}", result.best_loss);
println!("Stopped early: {}", result.stopped_early);
println!("Elapsed: {:.2}s", result.elapsed_secs);
}

train_epoch() - Single Epoch

Train for one epoch without callback overhead:

#![allow(unused)]
fn main() {
pub fn train_epoch<F, I>(&mut self, batches: I, forward_fn: F) -> f32
where
    F: Fn(&Tensor) -> Tensor,
    I: IntoIterator<Item = Batch>,
}

Returns: Average loss for the epoch

train_step() - Single Batch

Train on a single batch:

#![allow(unused)]
fn main() {
pub fn train_step<F>(&mut self, batch: &Batch, forward_fn: F) -> f32
where
    F: FnOnce(&Tensor) -> Tensor,
}

Returns: Loss for this batch

TrainResult

#![allow(unused)]
fn main() {
#[derive(Debug, Clone)]
pub struct TrainResult {
    pub final_epoch: usize,    // Last epoch completed
    pub final_loss: f32,       // Loss at final epoch
    pub best_loss: f32,        // Best loss achieved
    pub stopped_early: bool,   // Whether early stopping triggered
    pub elapsed_secs: f64,     // Total training time
}
}

Callback System

The trainer fires callbacks at six points in the training lifecycle:

CallbackAction

Callbacks return an action that controls training flow:

#![allow(unused)]
fn main() {
pub enum CallbackAction {
    Continue,   // Continue training normally
    Stop,       // Stop training immediately
    SkipEpoch,  // Skip to next epoch (epoch_begin only)
}
}

CallbackContext

Callbacks receive context with current training state:

#![allow(unused)]
fn main() {
pub struct CallbackContext {
    pub epoch: usize,           // Current epoch (0-indexed)
    pub max_epochs: usize,      // Maximum epochs
    pub step: usize,            // Current step in epoch
    pub steps_per_epoch: usize, // Total steps per epoch
    pub global_step: usize,     // Total steps across all epochs
    pub loss: f32,              // Current loss
    pub lr: f32,                // Current learning rate
    pub best_loss: Option<f32>, // Best loss so far
    pub val_loss: Option<f32>,  // Validation loss (if available)
    pub elapsed_secs: f64,      // Time since training started
}
}

Built-in Callbacks

EarlyStopping

Stop training when loss stops improving:

#![allow(unused)]
fn main() {
let es = EarlyStopping::new(
    5,      // patience: epochs without improvement
    0.001,  // min_delta: minimum improvement threshold
);
trainer.add_callback(es);
}

CheckpointCallback

Save model checkpoints:

#![allow(unused)]
fn main() {
let ckpt = CheckpointCallback::new("./checkpoints")
    .save_every(5)      // Save every 5 epochs
    .save_best(true);   // Also save best model
trainer.add_callback(ckpt);
}

ProgressCallback

Log training progress:

#![allow(unused)]
fn main() {
let progress = ProgressCallback::new(10);  // Log every 10 steps
trainer.add_callback(progress);
}

MonitorCallback

Real-time monitoring with NaN/Inf detection:

#![allow(unused)]
fn main() {
let monitor = MonitorCallback::new();
trainer.add_callback(monitor);
// Automatically stops training on NaN/Inf loss
}

Custom Callbacks

Implement TrainerCallback for custom behavior:

#![allow(unused)]
fn main() {
use entrenar::train::{TrainerCallback, CallbackContext, CallbackAction};

struct CustomCallback {
    // your state
}

impl TrainerCallback for CustomCallback {
    fn on_epoch_end(&mut self, ctx: &CallbackContext) -> CallbackAction {
        println!("Epoch {} complete, loss: {:.4}", ctx.epoch, ctx.loss);

        if ctx.loss > 100.0 {
            CallbackAction::Stop  // Loss exploded
        } else {
            CallbackAction::Continue
        }
    }

    fn name(&self) -> &str {
        "CustomCallback"
    }

    // Other methods have default implementations that return Continue
}

trainer.add_callback(CustomCallback { /* ... */ });
}

Accessing Trainer State

#![allow(unused)]
fn main() {
// Learning rate
let lr = trainer.lr();
trainer.set_lr(0.0001);

// Parameters
let params = trainer.params();
let params_mut = trainer.params_mut();

// Callbacks
let callbacks = trainer.callbacks();
let callbacks_mut = trainer.callbacks_mut();
}

Complete Example

use entrenar::train::{
    Trainer, TrainConfig, TrainResult, Batch, MSELoss,
    EarlyStopping, CheckpointCallback, ProgressCallback, MonitorCallback,
};
use entrenar::optim::Adam;
use entrenar::Tensor;

fn main() {
    // Model parameters
    let params = vec![
        Tensor::randn(784 * 256, true),  // Layer 1
        Tensor::randn(256 * 10, true),   // Layer 2
    ];

    // Optimizer
    let optimizer = Adam::new(0.001, 0.9, 0.999, 1e-8);

    // Config
    let config = TrainConfig::new()
        .with_max_grad_norm(1.0)
        .with_log_interval(100);

    // Create trainer
    let mut trainer = Trainer::new(params, Box::new(optimizer), config);
    trainer.set_loss(Box::new(MSELoss));

    // Add callbacks
    trainer.add_callback(EarlyStopping::new(10, 0.0001));
    trainer.add_callback(CheckpointCallback::new("./ckpt").save_every(5));
    trainer.add_callback(ProgressCallback::new(50));
    trainer.add_callback(MonitorCallback::new());

    // Training data
    let batches: Vec<Batch> = load_training_data();

    // Train
    let result: TrainResult = trainer.train(
        100,
        || batches.clone(),
        |x| forward_pass(x, trainer.params()),
    );

    // Results
    println!("Training complete!");
    println!("  Epochs: {}", result.final_epoch);
    println!("  Final loss: {:.6}", result.final_loss);
    println!("  Best loss: {:.6}", result.best_loss);
    println!("  Early stopped: {}", result.stopped_early);
    println!("  Time: {:.1}s", result.elapsed_secs);
}

See Also

• Train Config - Configuration options
• Early Stopping - Early stopping details
• Checkpointing - Checkpoint management
• Loss Functions - Available loss functions