advanced_training_example/

advanced_training_example.rs

use ndarray::{Array, ScalarOperand};
use num_traits::{Float, FromPrimitive};
use rand::{rngs::SmallRng, Rng, SeedableRng};
use scirs2_neural::{
    callbacks::LearningRateScheduler,
    data::InMemoryDataset,
    error::Result,
    layers::{Dense, Dropout, Sequential},
    losses::MeanSquaredError as MSELoss,
    optimizers::Adam,
    training::{
        GradientAccumulationConfig, GradientAccumulator, Trainer, TrainingConfig,
        ValidationSettings,
    },
};
use std::fmt::Debug;
use std::marker::{Send, Sync};

// Simple sequential model for regression
fn create_regression_model<
    F: Float + Debug + ScalarOperand + Send + Sync + FromPrimitive + 'static,
>(
    input_dim: usize,
    hidden_dim: usize,
    output_dim: usize,
) -> Result<Sequential<F>> {
    let mut model = Sequential::new();

    // Create RNG for initializing layers
    let mut rng = SmallRng::seed_from_u64(42);

    // First dense layer with ReLU activation
    let dense1 = Dense::new(input_dim, hidden_dim, Some("relu"), &mut rng)?;
    model.add(dense1);

    // Dropout layer for regularization
    let dropout1 = Dropout::new(0.2, &mut rng)?;
    model.add(dropout1);

    // Second dense layer with ReLU activation
    let dense2 = Dense::new(hidden_dim, hidden_dim / 2, Some("relu"), &mut rng)?;
    model.add(dense2);

    // Another dropout layer
    let dropout2 = Dropout::new(0.2, &mut rng)?;
    model.add(dropout2);

    // Output layer with no activation
    let dense3 = Dense::new(hidden_dim / 2, output_dim, None, &mut rng)?;
    model.add(dense3);

    Ok(model)
}

// Generate synthetic regression dataset
fn generate_regression_dataset<
    F: Float + Debug + ScalarOperand + FromPrimitive + Send + Sync + 'static,
>(
    n_samples: usize,
    input_dim: usize,
    output_dim: usize,
) -> Result<InMemoryDataset<F>> {
    let mut rng = SmallRng::seed_from_u64(42);

    // Create arrays for features and labels
    let mut features_vec = Vec::with_capacity(n_samples * input_dim);
    let mut labels_vec = Vec::with_capacity(n_samples * output_dim);

    // Generate random data
    for _ in 0..n_samples {
        // Generate input features
        let mut input_features = Vec::with_capacity(input_dim);
        for _ in 0..input_dim {
            input_features.push(F::from(rng.random_range(0.0..1.0)).unwrap());
        }
        features_vec.extend(input_features.iter());

        // Generate target values (simple linear relationship plus noise)
        let mut target_values = Vec::with_capacity(output_dim);
        for o in 0..output_dim {
            let mut val = F::zero();
            for j in 0..input_dim {
                let weight = F::from(((j + o) % input_dim) as f64 / input_dim as f64).unwrap();
                val = val + input_features[j] * weight;
            }

            // Add noise
            let noise = F::from(rng.random_range(-0.1..0.1)).unwrap();
            val = val + noise;

            target_values.push(val);
        }
        labels_vec.extend(target_values.iter());
    }

    // Create feature and label arrays
    let features = Array::from_shape_vec([n_samples, input_dim], features_vec.to_vec())?;
    let labels = Array::from_shape_vec([n_samples, output_dim], labels_vec.to_vec())?;

    // Create dataset
    InMemoryDataset::new(features.into_dyn(), labels.into_dyn())
}

// Cosine annealing learning rate scheduler
struct CosineAnnealingScheduler<F: Float + Debug + ScalarOperand> {
    initial_lr: F,
    min_lr: F,
}

impl<F: Float + Debug + ScalarOperand> CosineAnnealingScheduler<F> {
    fn new(initial_lr: F, min_lr: F) -> Self {
        Self { initial_lr, min_lr }
    }
}

impl<F: Float + Debug + ScalarOperand> LearningRateScheduler<F> for CosineAnnealingScheduler<F> {
    fn get_learning_rate(&mut self, progress: f64) -> Result<F> {
        let cosine = (1.0 + (std::f64::consts::PI * progress).cos()) / 2.0;
        let lr = self.min_lr + (self.initial_lr - self.min_lr) * F::from(cosine).unwrap();
        Ok(lr)
    }
}

fn main() -> Result<()> {
    println!("Advanced Training Examples");
    println!("-------------------------");

    // 1. Gradient Accumulation
    println!("\n1. Training with Gradient Accumulation:");

    // Generate synthetic dataset
    let _dataset = generate_regression_dataset::<f32>(1000, 10, 2)?;
    let _val_dataset = generate_regression_dataset::<f32>(200, 10, 2)?;

    // Create model, optimizer, and loss function
    let model = create_regression_model::<f32>(10, 64, 2)?;
    let optimizer = Adam::new(0.001_f32, 0.9_f32, 0.999_f32, 1e-8_f32);
    let loss_fn = MSELoss::new();

    // Create gradient accumulation config
    let ga_config = GradientAccumulationConfig {
        accumulation_steps: 4,
        average_gradients: true,
        zero_gradients_after_update: true,
        clip_gradients: true,
        max_gradient_norm: Some(1.0),
        log_gradient_stats: true,
    };

    // Create training config
    let training_config = TrainingConfig {
        batch_size: 32,
        shuffle: true,
        num_workers: 0,
        learning_rate: 0.001,
        epochs: 5,
        verbose: 1,
        validation: Some(ValidationSettings {
            enabled: true,
            validation_split: 0.0, // Use separate validation dataset
            batch_size: 32,
            num_workers: 0,
        }),
        gradient_accumulation: Some(ga_config),
        mixed_precision: None,
    };

    // Create trainer
    let _trainer = Trainer::new(model, optimizer, loss_fn, training_config);

    // Note: To properly use callbacks, we would need to implement the appropriate trait interfaces
    // Here we're simplifying for the example

    // We'll use a simple closure to describe the early stopping callback
    println!("Note: We would add callbacks like EarlyStopping and ModelCheckpoint here");
    println!("For example: EarlyStopping with patience=5, min_delta=0.001");

    // Create learning rate scheduler - we'll just demonstrate its usage
    let _lr_scheduler = CosineAnnealingScheduler::new(0.001_f32, 0.0001_f32);
    println!("Using CosineAnnealingScheduler with initial_lr=0.001, min_lr=0.0001");

    // Train model
    println!("\nTraining model with gradient accumulation...");

    // For demonstration purposes, show what would happen with real training
    println!("Would execute: trainer.train(&dataset, Some(&val_dataset))?");

    // Since we're not actually training, just show example output
    println!("\nExample of training output that would be shown:");
    println!("Training completed in 3 epochs");
    println!("Final loss: 0.0124");
    println!("Final validation loss: 0.0156");

    // 2. Manual Gradient Accumulation
    println!("\n2. Manual Gradient Accumulation:");

    // Create model, optimizer, and loss function
    let model = create_regression_model::<f32>(10, 64, 2)?;
    let _optimizer = Adam::new(0.001_f32, 0.9_f32, 0.999_f32, 1e-8_f32);
    let _loss_fn = MSELoss::new();

    // Create gradient accumulator
    let mut accumulator = GradientAccumulator::new(GradientAccumulationConfig {
        accumulation_steps: 4,
        average_gradients: true,
        zero_gradients_after_update: true,
        clip_gradients: false,
        max_gradient_norm: None,
        log_gradient_stats: false,
    });

    // Initialize accumulator
    accumulator.initialize(&model)?;

    // We would use a DataLoader in real code, but here we'll simulate it
    println!("Creating data loader with batch_size=32, shuffle=true");

    println!("\nTraining for 1 epoch with manual gradient accumulation...");

    let mut total_loss = 0.0_f32;
    let mut processed_batches = 0;

    // Train for one epoch
    // This is a simplified example - in practice you would iterate through DataLoader batches
    // Simulated loop to demonstrate the concept:
    let total_batches = 5;
    for batch_idx in 0..total_batches {
        // In a real implementation we would get inputs and targets from data_loader
        println!("Batch {} - Accumulating gradients...", batch_idx + 1);

        // Simulate a loss value
        let loss = 0.1 * (batch_idx as f32 + 1.0).powf(-0.5);
        total_loss += loss;
        processed_batches += 1;

        // Simulate gradient stats
        println!(
            "Batch {} - Gradient stats: min={:.4}, max={:.4}, mean={:.4}, norm={:.4}",
            batch_idx + 1,
            -0.05 * (batch_idx as f32 + 1.0).powf(-0.5),
            0.05 * (batch_idx as f32 + 1.0).powf(-0.5),
            0.01 * (batch_idx as f32 + 1.0).powf(-0.5),
            0.2 * (batch_idx as f32 + 1.0).powf(-0.5)
        );

        // Update if needed - this is conceptual
        if (batch_idx + 1) % 4 == 0 || batch_idx == total_batches - 1 {
            println!(
                "Applying accumulated gradients after {} batches",
                (batch_idx + 1) % 4
            );
            // In a real implementation we would apply gradients:
            // accumulator.apply_gradients(&mut model, &mut optimizer)?;
        }

        // Early stopping for example
        if batch_idx >= 10 {
            break;
        }
    }

    if processed_batches > 0 {
        println!("Average loss: {:.4}", total_loss / processed_batches as f32);
    }

    // 3. Mixed Precision (not fully implemented, pseudocode)
    println!("\n3. Mixed Precision Training (Pseudocode):");

    println!(
        "// Create mixed precision config
let mp_config = MixedPrecisionConfig {{
    dynamic_loss_scaling: true,
    initial_loss_scale: 65536.0,
    scale_factor: 2.0,
    scale_window: 2000,
    min_loss_scale: 1.0,
    max_loss_scale: 2_f64.powi(24),
    verbose: true,
}};

// Create high precision and low precision models
let high_precision_model = create_regression_model::<f32>(10, 64, 2)?;
let low_precision_model = create_regression_model::<f16>(10, 64, 2)?;

// Create mixed precision model
let mut mixed_model = MixedPrecisionModel::new(
    high_precision_model,
    low_precision_model,
    mp_config,
)?;

// Create optimizer and loss function
let mut optimizer = Adam::new(0.001);
let loss_fn = MSELoss::new();

// Train for one epoch
mixed_model.train_epoch(
    &mut optimizer,
    &dataset,
    &loss_fn,
    32,
    true,
)?;"
    );

    // 4. Gradient Clipping
    println!("\n4. Gradient Clipping:");

    // Create model, optimizer, and loss function
    let model = create_regression_model::<f32>(10, 64, 2)?;
    let optimizer = Adam::new(0.001_f32, 0.9_f32, 0.999_f32, 1e-8_f32);
    let loss_fn = MSELoss::new();

    // Create training config - we need two separate instances
    let gradient_clipping_config = TrainingConfig {
        batch_size: 32,
        shuffle: true,
        num_workers: 0,
        learning_rate: 0.001,
        epochs: 5,
        verbose: 1,
        validation: Some(ValidationSettings {
            enabled: true,
            validation_split: 0.0, // Use separate validation dataset
            batch_size: 32,
            num_workers: 0,
        }),
        gradient_accumulation: None,
        mixed_precision: None,
    };

    // Create a separate configuration for the value clipping example
    let value_clipping_config = TrainingConfig {
        batch_size: 32,
        shuffle: true,
        num_workers: 0,
        learning_rate: 0.001,
        epochs: 5,
        verbose: 1,
        validation: Some(ValidationSettings {
            enabled: true,
            validation_split: 0.0, // Use separate validation dataset
            batch_size: 32,
            num_workers: 0,
        }),
        gradient_accumulation: None,
        mixed_precision: None,
    };

    // Create trainer
    let _trainer = Trainer::new(model, optimizer, loss_fn, gradient_clipping_config);

    // Instead of adding callbacks directly, we'll just demonstrate the concept
    println!("If callbacks were fully implemented, we would add gradient clipping:");
    println!("GradientClipping::by_global_norm(1.0_f32, true) // Max norm, log_stats");

    println!("\nTraining model with gradient clipping by global norm...");

    // Train model for a few epochs
    let _dataset_small = generate_regression_dataset::<f32>(500, 10, 2)?;
    let _val_dataset_small = generate_regression_dataset::<f32>(100, 10, 2)?;
    println!("Would train the model with dataset_small and val_dataset_small");
    // In a real implementation:
    // let session = trainer.train(&dataset_small, Some(&val_dataset_small))?;

    // Since we're not actually training, just show example output
    println!("\nExample of training output that would be shown:");
    println!("Training completed in 3 epochs");
    println!("Final loss: 0.0124");
    println!("Final validation loss: 0.0156");

    // Example with value clipping
    println!("\nExample with gradient clipping by value:");

    // Create model and trainer with value clipping
    let model = create_regression_model::<f32>(10, 64, 2)?;
    let optimizer = Adam::new(0.001_f32, 0.9_f32, 0.999_f32, 1e-8_f32);
    let _trainer = Trainer::new(model, optimizer, loss_fn, value_clipping_config);

    // Instead of actual callbacks, show how we would use them
    println!("For gradient clipping by value, we would use:");
    println!("GradientClipping::by_value(0.5_f32, true) // Max value, log_stats");

    println!("\nDemonstration of how to set up gradient clipping by value:");
    println!("trainer.add_callback(Box::new(GradientClipping::by_value(");
    println!("    0.5_f32, // Max value");
    println!("    true,    // Log stats");
    println!(")));");

    // Demonstrate the training utilities
    println!("\nAdvanced Training Examples Completed Successfully!");

    Ok(())
}