Callback System

The callback system provides extensible hooks into the training loop, enabling behaviors like early stopping, checkpointing, progress logging, and real-time monitoring without modifying the core trainer.

Overview

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

// Add multiple callbacks
trainer.add_callback(EarlyStopping::new(5, 0.001));
trainer.add_callback(CheckpointCallback::new("./ckpt"));
trainer.add_callback(ProgressCallback::new(10));
trainer.add_callback(MonitorCallback::new());
}

Callback Lifecycle

Callbacks fire at six points during training:

train()
  │
  ├─► on_train_begin
  │
  ├─► for epoch in 0..max_epochs:
  │     │
  │     ├─► on_epoch_begin
  │     │
  │     ├─► for batch in batches:
  │     │     ├─► on_step_begin
  │     │     ├─► train_step()
  │     │     └─► on_step_end
  │     │
  │     └─► on_epoch_end
  │
  └─► on_train_end

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)
}
}

Behavior:

Continue - Training proceeds normally
Stop - Training stops, TrainResult.stopped_early = true
SkipEpoch - Skip remaining steps in current epoch (only valid in on_epoch_begin)

CallbackContext

Every callback receives 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 configured
    pub step: usize,            // Current step within epoch
    pub steps_per_epoch: usize, // Total steps in epoch
    pub global_step: usize,     // Total steps across all epochs
    pub loss: f32,              // Current/latest loss
    pub lr: f32,                // Current learning rate
    pub best_loss: Option<f32>, // Best loss achieved so far
    pub val_loss: Option<f32>,  // Validation loss (if provided)
    pub elapsed_secs: f64,      // Seconds since training started
}
}

TrainerCallback Trait

#![allow(unused)]
fn main() {
pub trait TrainerCallback: Send {
    /// Called before training begins
    fn on_train_begin(&mut self, ctx: &CallbackContext) -> CallbackAction {
        CallbackAction::Continue
    }

    /// Called after training ends
    fn on_train_end(&mut self, ctx: &CallbackContext) {}

    /// Called at the start of each epoch
    fn on_epoch_begin(&mut self, ctx: &CallbackContext) -> CallbackAction {
        CallbackAction::Continue
    }

    /// Called at the end of each epoch
    fn on_epoch_end(&mut self, ctx: &CallbackContext) -> CallbackAction {
        CallbackAction::Continue
    }

    /// Called before each training step
    fn on_step_begin(&mut self, ctx: &CallbackContext) -> CallbackAction {
        CallbackAction::Continue
    }

    /// Called after each training step
    fn on_step_end(&mut self, ctx: &CallbackContext) -> CallbackAction {
        CallbackAction::Continue
    }

    /// Callback name for logging
    fn name(&self) -> &str;
}
}

Built-in Callbacks

EarlyStopping

Stops training when loss stops improving:

#![allow(unused)]
fn main() {
pub struct EarlyStopping {
    patience: usize,      // Epochs to wait before stopping
    min_delta: f32,       // Minimum improvement threshold
    best_loss: f32,       // Best loss seen
    epochs_without_improvement: usize,
}

// Usage
let es = EarlyStopping::new(5, 0.001);
// Stops if loss doesn't improve by at least 0.001 for 5 epochs
}

Behavior:

Tracks best loss seen during training
Counts epochs without improvement (loss not decreasing by min_delta)
Returns CallbackAction::Stop when patience exhausted

CheckpointCallback

Saves model checkpoints periodically and/or when best loss achieved:

#![allow(unused)]
fn main() {
pub struct CheckpointCallback {
    save_dir: PathBuf,
    save_every: Option<usize>,  // Save every N epochs
    save_best: bool,            // Save when best loss achieved
    best_loss: f32,
}

// Usage
let ckpt = CheckpointCallback::new("./checkpoints")
    .save_every(5)      // Save every 5 epochs
    .save_best(true);   // Also save best model

// Creates files like:
// ./checkpoints/checkpoint_epoch_5.json
// ./checkpoints/checkpoint_epoch_10.json
// ./checkpoints/checkpoint_best.json
}

ProgressCallback

Logs training progress to stdout:

#![allow(unused)]
fn main() {
pub struct ProgressCallback {
    log_interval: usize,  // Log every N steps
}

// Usage
let progress = ProgressCallback::new(10);
// Logs: "Epoch 1/100 [========>  ] Step 50/500: loss: 0.1234"
}

Output format:

Epoch 1/100 [========>  ] loss: 0.2345, lr: 0.001, elapsed: 12.3s
  Step 10/100: loss: 0.2456
  Step 20/100: loss: 0.2234

MonitorCallback

Real-time monitoring with NaN/Inf detection and metrics collection:

#![allow(unused)]
fn main() {
pub struct MonitorCallback {
    collector: MetricsCollector,
    andon: AndonSystem,
}

// Usage
let monitor = MonitorCallback::new();
// Automatically:
// - Records loss, learning rate metrics
// - Detects NaN/Inf and triggers Stop
// - Integrates with Andon alerting system
}

Automatic detection:

NaN loss → CallbackAction::Stop
Inf loss → CallbackAction::Stop
Triggers Andon alert for investigation

Custom Callbacks

Basic Example

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

struct LossLogger {
    losses: Vec<f32>,
}

impl LossLogger {
    fn new() -> Self {
        Self { losses: Vec::new() }
    }
}

impl TrainerCallback for LossLogger {
    fn on_epoch_end(&mut self, ctx: &CallbackContext) -> CallbackAction {
        self.losses.push(ctx.loss);
        println!("Epoch {}: loss = {:.6}", ctx.epoch, ctx.loss);
        CallbackAction::Continue
    }

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

Learning Rate Warmup

#![allow(unused)]
fn main() {
struct WarmupCallback {
    warmup_epochs: usize,
    target_lr: f32,
}

impl TrainerCallback for WarmupCallback {
    fn on_epoch_begin(&mut self, ctx: &CallbackContext) -> CallbackAction {
        if ctx.epoch < self.warmup_epochs {
            let warmup_lr = self.target_lr * (ctx.epoch + 1) as f32
                          / self.warmup_epochs as f32;
            // Would need trainer access to set LR
            println!("Warmup LR: {:.6}", warmup_lr);
        }
        CallbackAction::Continue
    }

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

Gradient Explosion Detector

#![allow(unused)]
fn main() {
struct GradientMonitor {
    max_loss: f32,
    loss_history: Vec<f32>,
}

impl TrainerCallback for GradientMonitor {
    fn on_step_end(&mut self, ctx: &CallbackContext) -> CallbackAction {
        self.loss_history.push(ctx.loss);

        // Detect sudden loss spike
        if self.loss_history.len() > 1 {
            let prev = self.loss_history[self.loss_history.len() - 2];
            if ctx.loss > prev * 10.0 {
                eprintln!("WARNING: Loss spike detected! {} -> {}", prev, ctx.loss);
                return CallbackAction::Stop;
            }
        }

        if ctx.loss > self.max_loss {
            eprintln!("ERROR: Loss exceeded threshold: {} > {}", ctx.loss, self.max_loss);
            return CallbackAction::Stop;
        }

        CallbackAction::Continue
    }

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

CallbackManager

The CallbackManager orchestrates multiple callbacks:

#![allow(unused)]
fn main() {
pub struct CallbackManager {
    callbacks: Vec<Box<dyn TrainerCallback>>,
}

impl CallbackManager {
    pub fn new() -> Self;
    pub fn add<C: TrainerCallback + 'static>(&mut self, callback: C);
    pub fn is_empty(&self) -> bool;
    pub fn len(&self) -> usize;

    // Event dispatchers (called by Trainer)
    pub fn on_train_begin(&mut self, ctx: &CallbackContext) -> CallbackAction;
    pub fn on_train_end(&mut self, ctx: &CallbackContext);
    pub fn on_epoch_begin(&mut self, ctx: &CallbackContext) -> CallbackAction;
    pub fn on_epoch_end(&mut self, ctx: &CallbackContext) -> CallbackAction;
    pub fn on_step_begin(&mut self, ctx: &CallbackContext) -> CallbackAction;
    pub fn on_step_end(&mut self, ctx: &CallbackContext) -> CallbackAction;
}
}

Dispatch behavior:

Callbacks fire in order they were added
If any callback returns Stop, remaining callbacks don't fire
on_train_end always fires (even after early stop)

Best Practices

Callback Order

Add callbacks in order of priority:

#![allow(unused)]
fn main() {
// Critical monitoring first
trainer.add_callback(MonitorCallback::new());     // NaN detection
trainer.add_callback(EarlyStopping::new(5, 0.001)); // Early stopping

// Logging/checkpointing after
trainer.add_callback(ProgressCallback::new(10));
trainer.add_callback(CheckpointCallback::new("./ckpt"));
}

Stateful Callbacks

Callbacks can maintain state across training:

#![allow(unused)]
fn main() {
struct StatefulCallback {
    epoch_losses: Vec<f32>,
    best_epoch: usize,
}

impl TrainerCallback for StatefulCallback {
    fn on_epoch_end(&mut self, ctx: &CallbackContext) -> CallbackAction {
        self.epoch_losses.push(ctx.loss);

        if ctx.best_loss == Some(ctx.loss) {
            self.best_epoch = ctx.epoch;
        }

        CallbackAction::Continue
    }

    fn on_train_end(&mut self, ctx: &CallbackContext) {
        println!("Best epoch: {} with loss {:.6}",
            self.best_epoch,
            self.epoch_losses[self.best_epoch]);
    }

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

Thread Safety

Callbacks must be Send to support potential future parallelism:

#![allow(unused)]
fn main() {
// Good: Uses Arc for shared state
struct ThreadSafeCallback {
    counter: Arc<AtomicUsize>,
}

// Bad: Uses Rc (not Send)
struct NotSendCallback {
    counter: Rc<RefCell<usize>>,  // Won't compile!
}
}

Entrenar - Training & Optimization Library