Struct Trainer 

pub struct Trainer { /* private fields */ }

Training utilities with scheduler, metrics, and validation

Implementations

impl Trainer

pub fn new(model: TrainableSSM, config: TrainingConfig) -> CoreResult<Self>

Create a new trainer

pub fn train_epoch<F>( &mut self, data_loader: &[(Tensor, Tensor)], loss_fn: F, ) -> CoreResult<f32>

where F: Fn(&Tensor, &Tensor) -> CoreResult<Tensor>,

Train for one epoch

pub fn evaluate<F>( &self, data_loader: &[(Tensor, Tensor)], loss_fn: F, ) -> CoreResult<f32>

where F: Fn(&Tensor, &Tensor) -> CoreResult<Tensor>,

Evaluate on validation data

pub fn fit<F>( &mut self, train_loader: TimeSeriesDataLoader, val_loader: Option<TimeSeriesDataLoader>, loss_fn: F, ) -> CoreResult<()>

where F: Fn(&Tensor, &Tensor) -> CoreResult<Tensor> + Copy,

Full training loop with validation and early stopping

pub fn model(&self) -> &TrainableSSM

Get reference to the model

pub fn model_mut(&mut self) -> &mut TrainableSSM

Get mutable reference to the model

pub fn metrics(&self) -> &TrainingMetrics

Get reference to training metrics

pub fn metrics_mut(&mut self) -> &mut TrainingMetrics

Get mutable reference to training metrics

pub fn current_step(&self) -> usize

Get current training step

pub fn save_checkpoint<P: AsRef<Path>>( &self, path: P, name: &str, ) -> CoreResult<()>

Save checkpoint to disk

Saves model weights, optimizer state, training configuration, metrics, and metadata.

Arguments

• path - Directory to save checkpoint files
• name - Checkpoint name (without extension)

Example

trainer.save_checkpoint("checkpoints", "epoch_10")?;
// Creates: checkpoints/epoch_10.safetensors and checkpoints/epoch_10.json

pub fn load_checkpoint<P: AsRef<Path>>( path: P, name: &str, model_config: KizzasiConfig, ) -> CoreResult<Self>

Load checkpoint and resume training

Creates a new Trainer from a saved checkpoint, restoring model weights, configuration, and training state.

Arguments

• path - Directory containing checkpoint files
• name - Checkpoint name (without extension)
• model_config - Model configuration (must match saved model)

Example

let trainer = Trainer::load_checkpoint("checkpoints", "epoch_10", model_config)?;
// Continue training from epoch 10

pub fn save_checkpoint_auto<P: AsRef<Path>>(&self, path: P) -> CoreResult<()>

Save checkpoint with automatic naming (epoch-based)

Convenience method that automatically names checkpoints based on current epoch.

Example

trainer.save_checkpoint_auto("checkpoints")?;
// Creates: checkpoints/checkpoint_epoch_5.safetensors, etc.

pub fn save_best_checkpoint<P: AsRef<Path>>(&self, path: P) -> CoreResult<()>

Save checkpoint if this is the best epoch (lowest validation loss)

Automatically saves a “best” checkpoint when validation loss improves.

Example

// After each validation epoch
trainer.save_best_checkpoint("checkpoints")?;