Struct Trainer 

pub struct Trainer {
    pub metrics: MetricsTracker,
    /* private fields */
}

High-level trainer that orchestrates the training loop

Example

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

// Setup
let params = vec![Tensor::zeros(10, 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));
trainer.add_callback(EarlyStopping::new(5, 0.001));

// Training with callbacks
// let result = trainer.train(10, || batches.clone(), |x| x.clone());

Fields

metrics: MetricsTracker

Metrics tracker

Implementations

impl Trainer

pub fn new( params: Vec<Tensor>, optimizer: Box<dyn Optimizer>, config: TrainConfig, ) -> Self

Create a new trainer

pub fn set_loss(&mut self, loss_fn: Box<dyn LossFn>)

Set the loss function

pub fn add_callback<C: TrainerCallback + 'static>(&mut self, callback: C)

Add a callback to the trainer

pub fn lr(&self) -> f32

Get current learning rate

pub fn set_lr(&mut self, lr: f32)

Set learning rate

pub fn params(&self) -> &[Tensor]

Get reference to model parameters

pub fn params_mut(&mut self) -> &mut [Tensor]

Get mutable reference to model parameters

pub fn callbacks(&self) -> &CallbackManager

Get reference to callback manager

pub fn callbacks_mut(&mut self) -> &mut CallbackManager

Get mutable reference to callback manager

pub fn save( &self, path: impl AsRef<Path>, name: &str, architecture: &str, ) -> Result<()>

Save model parameters to a file

This method persists the trained model weights to disk in SafeTensors format. Call this after training completes to preserve the learned parameters.

Arguments

• path - Output file path (should end in .safetensors)
• name - Model name for metadata
• architecture - Model architecture description

Example

// After training...
trainer.save("model.safetensors", "my-model", "linear").expect("save failed");

Errors

Returns an error if the file cannot be written.

pub fn save_with_names( &self, path: impl AsRef<Path>, name: &str, architecture: &str, param_names: &[&str], ) -> Result<()>

Save model with custom parameter names

Like save() but allows specifying custom names for each parameter tensor.

Arguments

• path - Output file path
• name - Model name
• architecture - Architecture description
• param_names - Names for each parameter (must match params length)

Errors

Returns an error if param_names length doesn’t match params or file cannot be written.

impl Trainer

pub fn train_epoch<F, I>(&mut self, batches: I, forward_fn: F) -> f32

where F: Fn(&Tensor) -> Tensor, I: IntoIterator<Item = Batch>,

Train for one epoch

Arguments

• batches - Iterator over training batches
• forward_fn - Closure that computes predictions from inputs

Returns

Average loss over the epoch

pub fn validate<F, I>(&mut self, batches: I, forward_fn: F) -> f32

where F: Fn(&Tensor) -> Tensor, I: IntoIterator<Item = Batch>,

Validate on a dataset without updating parameters

Arguments

• batches - Iterator over validation batches
• forward_fn - Closure that computes predictions from inputs

Returns

Average validation loss

Example

let val_loss = trainer.validate(val_batches, |x| x.clone());
println!("Validation loss: {:.4}", val_loss);

impl Trainer

pub fn train_step<F>(&mut self, batch: &Batch, forward_fn: F) -> f32

where F: FnOnce(&Tensor) -> Tensor,

Perform a single training step

Arguments

• batch - Training batch with inputs and targets
• forward_fn - Closure that computes predictions from inputs

Returns

Scalar loss value for this batch

Example

let loss = trainer.train_step(&batch, |inputs| {
    // Forward pass: compute predictions
    inputs.clone() // Simplified example
});

impl Trainer

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>,

Train for multiple epochs with full callback support

Arguments

• max_epochs - Maximum number of epochs to train
• batch_fn - Function that returns batches for each epoch
• forward_fn - Closure that computes predictions from inputs

Returns

TrainResult with final metrics

Example

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

let result = trainer.train(100, || batches.clone(), |x| x.clone());
println!("Trained {} epochs, final loss: {:.4}", result.final_epoch, result.final_loss);

impl Trainer

pub fn train_with_val<F, BT, BV, IT, IV>( &mut self, max_epochs: usize, train_fn: BT, val_fn: BV, forward_fn: F, ) -> TrainResult

where F: Fn(&Tensor) -> Tensor, BT: Fn() -> IT, BV: Fn() -> IV, IT: IntoIterator<Item = Batch>, IV: IntoIterator<Item = Batch>,

Train for multiple epochs with validation after each epoch

This method runs training and validation each epoch, passing validation loss to callbacks for proper early stopping and checkpointing.

Arguments

• max_epochs - Maximum number of epochs to train
• train_fn - Function that returns training batches for each epoch
• val_fn - Function that returns validation batches for each epoch
• forward_fn - Closure that computes predictions from inputs

Returns

TrainResult with final metrics including best validation loss

Example

trainer.add_callback(EarlyStopping::new(5, 0.001).monitor_validation());

let result = trainer.train_with_val(
    100,
    || train_batches.clone(),
    || val_batches.clone(),
    |x| x.clone()
);
println!("Best val loss: {:.4}", result.best_loss);