osobhandClaude Opus 4.6
02d382d5f6
Consistent formatting pass: line wrapping, import sorting, trailing whitespace removal, let-chain indentation, merged derive attributes, and unsafe block reformatting. Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
2406 lines
83 KiB
Rust
2406 lines
83 KiB
Rust
//! FNO Training module for neural operator demo
|
|
//!
|
|
//! Provides training functionality for Fourier Neural Operators using
|
|
//! pure Rust implementation with rtx-autograd and rtx-losses.
|
|
//!
|
|
//! # Architecture
|
|
//!
|
|
//! The trainer uses a simplified FNO architecture suitable for training:
|
|
//! - Trainable parameters stored as flat vectors
|
|
//! - Forward pass through `FNO2d` model
|
|
//! - MSE loss computation
|
|
//! - `AdamW` optimizer for parameter updates
|
|
//!
|
|
//! # Example
|
|
//!
|
|
//! ```rust,ignore
|
|
//! use rtx_neural_operator_demo::training::{FnoTrainer, FnoTrainingConfig};
|
|
//!
|
|
//! let config = FnoTrainingConfig::standard();
|
|
//! let trainer = FnoTrainer::new(config);
|
|
//!
|
|
//! // Start training (runs in background)
|
|
//! trainer.start_training().await?;
|
|
//!
|
|
//! // Poll for progress
|
|
//! let progress = trainer.get_progress();
|
|
//! ```
|
|
|
|
use std::sync::Arc;
|
|
use std::sync::atomic::{AtomicBool, Ordering};
|
|
use std::time::Instant;
|
|
|
|
use rtx_backend::Backend;
|
|
use rtx_backend::auto_select::{BackendType, get_backend_info, select_best_backend};
|
|
use rtx_backend_cpu::CpuBackend;
|
|
use rtx_neural_operator::FNO2d;
|
|
use rtx_neural_operator::weights::FNO2dWeights;
|
|
use rtx_neural_operator_shared::config::PDEConfig;
|
|
use rtx_neural_operator_shared::ipc::{TrainingConfig, TrainingProgress, TrainingStatus};
|
|
use rtx_nn::GenericModule4D;
|
|
use rtx_tensor::GenericTensor;
|
|
use tokio::sync::broadcast;
|
|
|
|
use crate::data::{DarcyDataConfig, DarcyDataGenerator, DarcySample, TrainingBatch};
|
|
|
|
// Conditional GPU backend imports
|
|
#[cfg(feature = "cuda")]
|
|
use rtx_backend_cuda::{CudaBackend, CudaDevice};
|
|
#[cfg(feature = "metal")]
|
|
use rtx_backend_metal::{MetalBackend, MetalDevice};
|
|
#[cfg(feature = "rocm")]
|
|
use rtx_backend_rocm::{RocmBackend, RocmDevice};
|
|
|
|
/// Get a human-readable device string for a specific backend type.
|
|
fn get_device_string_for_backend(backend_type: BackendType) -> String {
|
|
let info = get_backend_info(backend_type);
|
|
|
|
match backend_type {
|
|
BackendType::Cpu => {
|
|
if let Some(ref device_name) = info.device_name {
|
|
format!(
|
|
"CPU ({})",
|
|
device_name
|
|
.split_whitespace()
|
|
.take(4)
|
|
.collect::<Vec<_>>()
|
|
.join(" ")
|
|
)
|
|
} else {
|
|
"CPU".to_string()
|
|
}
|
|
}
|
|
BackendType::Cuda => {
|
|
if let Some(ref device_name) = info.device_name {
|
|
format!("CUDA ({device_name})")
|
|
} else {
|
|
"CUDA".to_string()
|
|
}
|
|
}
|
|
BackendType::Metal => {
|
|
if let Some(ref device_name) = info.device_name {
|
|
format!("Metal ({device_name})")
|
|
} else {
|
|
"Metal".to_string()
|
|
}
|
|
}
|
|
BackendType::Rocm => {
|
|
if let Some(ref device_name) = info.device_name {
|
|
format!("ROCm ({device_name})")
|
|
} else {
|
|
"ROCm".to_string()
|
|
}
|
|
}
|
|
_ => backend_type.to_string(),
|
|
}
|
|
}
|
|
|
|
/// Get a human-readable device string for the current best available backend.
|
|
fn get_device_string() -> String {
|
|
let backend_type = select_best_backend();
|
|
get_device_string_for_backend(backend_type)
|
|
}
|
|
|
|
// ============================================================================
|
|
// AdamW Optimizer
|
|
// ============================================================================
|
|
|
|
/// `AdamW` optimizer with decoupled weight decay.
|
|
///
|
|
/// Implements the `AdamW` algorithm for parameter optimization with:
|
|
/// - Adaptive learning rates per parameter
|
|
/// - Momentum (first moment) tracking
|
|
/// - Variance (second moment) tracking
|
|
/// - Decoupled weight decay for better generalization
|
|
#[derive(Debug)]
|
|
struct AdamW {
|
|
/// First moment estimates (momentum)
|
|
m: Vec<Vec<f32>>,
|
|
/// Second moment estimates (variance)
|
|
v: Vec<Vec<f32>>,
|
|
/// Timestep counter for bias correction
|
|
t: u64,
|
|
/// Learning rate
|
|
lr: f32,
|
|
/// First moment decay rate (β₁)
|
|
beta1: f32,
|
|
/// Second moment decay rate (β₂)
|
|
beta2: f32,
|
|
/// Numerical stability constant (ε)
|
|
eps: f32,
|
|
/// Weight decay coefficient (λ)
|
|
weight_decay: f32,
|
|
}
|
|
|
|
impl AdamW {
|
|
/// Create a new `AdamW` optimizer.
|
|
///
|
|
/// # Arguments
|
|
/// * `param_sizes` - Size of each parameter vector
|
|
/// * `lr` - Learning rate (typical: 1e-3)
|
|
/// * `beta1` - First moment decay (typical: 0.9)
|
|
/// * `beta2` - Second moment decay (typical: 0.999)
|
|
/// * `weight_decay` - Weight decay coefficient (typical: 0.01)
|
|
fn new(param_sizes: &[usize], lr: f32, beta1: f32, beta2: f32, weight_decay: f32) -> Self {
|
|
let m = param_sizes.iter().map(|&size| vec![0.0f32; size]).collect();
|
|
let v = param_sizes.iter().map(|&size| vec![0.0f32; size]).collect();
|
|
|
|
Self {
|
|
m,
|
|
v,
|
|
t: 0,
|
|
lr,
|
|
beta1,
|
|
beta2,
|
|
eps: 1e-8,
|
|
weight_decay,
|
|
}
|
|
}
|
|
|
|
/// Create `AdamW` with default hyperparameters.
|
|
///
|
|
/// Uses lr=1e-3, β₁=0.9, β₂=0.999, `weight_decay=0.01`
|
|
#[allow(dead_code)]
|
|
fn with_defaults(param_sizes: &[usize]) -> Self {
|
|
Self::new(param_sizes, 1e-3, 0.9, 0.999, 0.01)
|
|
}
|
|
|
|
/// Perform one optimization step.
|
|
///
|
|
/// Updates parameters in-place using `AdamW` update rule:
|
|
/// ```text
|
|
/// m = β₁ * m + (1 - β₁) * g // Update momentum
|
|
/// v = β₂ * v + (1 - β₂) * g² // Update variance
|
|
/// m̂ = m / (1 - β₁^t) // Bias-corrected momentum
|
|
/// v̂ = v / (1 - β₂^t) // Bias-corrected variance
|
|
/// θ = θ - lr * (m̂ / (√v̂ + ε) + λ * θ) // Decoupled weight decay update
|
|
/// ```
|
|
fn step(&mut self, params: &mut [Vec<f32>], grads: &[Vec<f32>]) {
|
|
self.t += 1;
|
|
|
|
// Bias correction factors
|
|
let bias_correction1 = 1.0 - self.beta1.powi(self.t as i32);
|
|
let bias_correction2 = 1.0 - self.beta2.powi(self.t as i32);
|
|
|
|
for (i, (param, grad)) in params.iter_mut().zip(grads.iter()).enumerate() {
|
|
for (j, (p, &g)) in param.iter_mut().zip(grad.iter()).enumerate() {
|
|
// Update first moment: m = β₁ * m + (1 - β₁) * g
|
|
self.m[i][j] = self.beta1 * self.m[i][j] + (1.0 - self.beta1) * g;
|
|
|
|
// Update second moment: v = β₂ * v + (1 - β₂) * g²
|
|
self.v[i][j] = self.beta2 * self.v[i][j] + (1.0 - self.beta2) * g * g;
|
|
|
|
// Bias-corrected estimates
|
|
let m_hat = self.m[i][j] / bias_correction1;
|
|
let v_hat = self.v[i][j] / bias_correction2;
|
|
|
|
// AdamW update with decoupled weight decay
|
|
let adam_update = m_hat / (v_hat.sqrt() + self.eps);
|
|
*p -= self.lr * (adam_update + self.weight_decay * *p);
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Get current learning rate
|
|
#[allow(dead_code)]
|
|
fn learning_rate(&self) -> f32 {
|
|
self.lr
|
|
}
|
|
|
|
/// Set learning rate
|
|
#[allow(dead_code)]
|
|
fn set_learning_rate(&mut self, lr: f32) {
|
|
self.lr = lr;
|
|
}
|
|
}
|
|
|
|
/// Extract all trainable parameters from `FNO2dWeights` as flat vectors.
|
|
fn extract_params(weights: &FNO2dWeights) -> Vec<Vec<f32>> {
|
|
let mut params = Vec::new();
|
|
|
|
// Lifting MLP parameters
|
|
params.push(weights.lifting_fc1_weight.clone());
|
|
params.push(weights.lifting_fc1_bias.clone());
|
|
params.push(weights.lifting_fc2_weight.clone());
|
|
params.push(weights.lifting_fc2_bias.clone());
|
|
|
|
// Spectral conv parameters (for each layer)
|
|
for sw in &weights.spectral_weights {
|
|
params.push(sw.weights1_real.clone());
|
|
params.push(sw.weights1_imag.clone());
|
|
params.push(sw.weights2_real.clone());
|
|
params.push(sw.weights2_imag.clone());
|
|
}
|
|
|
|
// 1x1 conv parameters
|
|
for (weight, bias) in &weights.conv_weights {
|
|
params.push(weight.clone());
|
|
params.push(bias.clone());
|
|
}
|
|
|
|
// Projection parameters
|
|
for (weight, bias) in &weights.projection_weights {
|
|
params.push(weight.clone());
|
|
params.push(bias.clone());
|
|
}
|
|
|
|
params
|
|
}
|
|
|
|
/// Update `FNO2dWeights` with new parameter values.
|
|
fn update_weights(weights: &mut FNO2dWeights, params: &[Vec<f32>]) {
|
|
let mut idx = 0;
|
|
|
|
// Lifting MLP parameters
|
|
weights.lifting_fc1_weight = params[idx].clone();
|
|
idx += 1;
|
|
weights.lifting_fc1_bias = params[idx].clone();
|
|
idx += 1;
|
|
weights.lifting_fc2_weight = params[idx].clone();
|
|
idx += 1;
|
|
weights.lifting_fc2_bias = params[idx].clone();
|
|
idx += 1;
|
|
|
|
// Spectral conv parameters
|
|
for sw in &mut weights.spectral_weights {
|
|
sw.weights1_real = params[idx].clone();
|
|
idx += 1;
|
|
sw.weights1_imag = params[idx].clone();
|
|
idx += 1;
|
|
sw.weights2_real = params[idx].clone();
|
|
idx += 1;
|
|
sw.weights2_imag = params[idx].clone();
|
|
idx += 1;
|
|
}
|
|
|
|
// 1x1 conv parameters
|
|
for (weight, bias) in &mut weights.conv_weights {
|
|
*weight = params[idx].clone();
|
|
idx += 1;
|
|
*bias = params[idx].clone();
|
|
idx += 1;
|
|
}
|
|
|
|
// Projection parameters
|
|
for (weight, bias) in &mut weights.projection_weights {
|
|
*weight = params[idx].clone();
|
|
idx += 1;
|
|
*bias = params[idx].clone();
|
|
idx += 1;
|
|
}
|
|
}
|
|
|
|
// ============================================================================
|
|
// Early Stopping
|
|
// ============================================================================
|
|
|
|
/// Early stopping monitor to prevent overfitting.
|
|
///
|
|
/// Tracks validation loss and triggers early stopping when no improvement
|
|
/// is seen for a specified number of epochs (patience).
|
|
#[derive(Debug)]
|
|
struct EarlyStopping {
|
|
/// Number of epochs to wait before stopping
|
|
patience: usize,
|
|
/// Minimum improvement required to reset counter
|
|
min_delta: f32,
|
|
/// Best validation loss seen so far
|
|
best_loss: f32,
|
|
/// Epochs since last improvement
|
|
epochs_without_improvement: usize,
|
|
}
|
|
|
|
impl EarlyStopping {
|
|
/// Create a new early stopping monitor.
|
|
///
|
|
/// # Arguments
|
|
/// * `patience` - Number of epochs to wait for improvement
|
|
/// * `min_delta` - Minimum improvement required (default: 1e-4)
|
|
fn new(patience: usize, min_delta: f32) -> Self {
|
|
Self {
|
|
patience,
|
|
min_delta,
|
|
best_loss: f32::INFINITY,
|
|
epochs_without_improvement: 0,
|
|
}
|
|
}
|
|
|
|
/// Check if training should stop.
|
|
///
|
|
/// Returns `true` if no improvement has been seen for `patience` epochs.
|
|
fn check(&mut self, val_loss: f32) -> bool {
|
|
if val_loss < self.best_loss - self.min_delta {
|
|
// Improvement found
|
|
self.best_loss = val_loss;
|
|
self.epochs_without_improvement = 0;
|
|
false
|
|
} else {
|
|
// No improvement
|
|
self.epochs_without_improvement += 1;
|
|
self.epochs_without_improvement >= self.patience
|
|
}
|
|
}
|
|
|
|
/// Get the best loss seen so far
|
|
#[allow(dead_code)]
|
|
fn best_loss(&self) -> f32 {
|
|
self.best_loss
|
|
}
|
|
|
|
/// Get epochs without improvement
|
|
fn epochs_without_improvement(&self) -> usize {
|
|
self.epochs_without_improvement
|
|
}
|
|
}
|
|
|
|
// ============================================================================
|
|
// Learning Rate Scheduler
|
|
// ============================================================================
|
|
|
|
/// Learning rate scheduling strategy.
|
|
#[derive(Debug, Clone)]
|
|
pub enum LRSchedulerType {
|
|
/// Reduce LR when validation loss plateaus
|
|
ReduceOnPlateau {
|
|
/// Factor to reduce LR by (e.g., 0.1 = reduce to 10%)
|
|
factor: f32,
|
|
/// Epochs to wait before reducing
|
|
patience: usize,
|
|
/// Minimum learning rate
|
|
min_lr: f32,
|
|
},
|
|
/// Cosine annealing from initial LR to min LR
|
|
CosineAnnealing {
|
|
/// Total epochs for one cycle
|
|
t_max: usize,
|
|
/// Minimum learning rate
|
|
eta_min: f32,
|
|
},
|
|
/// Step decay: reduce LR every N epochs
|
|
StepDecay {
|
|
/// Reduce LR every `step_size` epochs
|
|
step_size: usize,
|
|
/// Factor to multiply LR by
|
|
gamma: f32,
|
|
},
|
|
}
|
|
|
|
/// Learning rate scheduler.
|
|
#[derive(Debug)]
|
|
pub struct LRScheduler {
|
|
/// Scheduling strategy
|
|
scheduler_type: LRSchedulerType,
|
|
/// Initial learning rate
|
|
initial_lr: f32,
|
|
/// Current learning rate
|
|
current_lr: f32,
|
|
/// Best validation loss (for `ReduceOnPlateau`)
|
|
best_val_loss: f32,
|
|
/// Epochs since improvement (for `ReduceOnPlateau`)
|
|
epochs_since_improvement: usize,
|
|
}
|
|
|
|
impl LRScheduler {
|
|
/// Create a new LR scheduler.
|
|
#[must_use]
|
|
pub fn new(scheduler_type: LRSchedulerType, initial_lr: f32) -> Self {
|
|
Self {
|
|
scheduler_type,
|
|
initial_lr,
|
|
current_lr: initial_lr,
|
|
best_val_loss: f32::INFINITY,
|
|
epochs_since_improvement: 0,
|
|
}
|
|
}
|
|
|
|
/// Update learning rate based on current epoch and validation loss.
|
|
///
|
|
/// Returns the new learning rate.
|
|
pub fn step(&mut self, epoch: usize, val_loss: Option<f32>) -> f32 {
|
|
match &self.scheduler_type {
|
|
LRSchedulerType::ReduceOnPlateau {
|
|
factor,
|
|
patience,
|
|
min_lr,
|
|
} => {
|
|
if let Some(loss) = val_loss {
|
|
if loss < self.best_val_loss {
|
|
self.best_val_loss = loss;
|
|
self.epochs_since_improvement = 0;
|
|
} else {
|
|
self.epochs_since_improvement += 1;
|
|
if self.epochs_since_improvement >= *patience {
|
|
self.current_lr = (self.current_lr * factor).max(*min_lr);
|
|
self.epochs_since_improvement = 0;
|
|
tracing::info!("LR reduced to {:.2e}", self.current_lr);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
LRSchedulerType::CosineAnnealing { t_max, eta_min } => {
|
|
// Cosine annealing: lr = eta_min + 0.5 * (initial_lr - eta_min) * (1 + cos(pi * t / T))
|
|
let t = epoch % t_max;
|
|
let cos_term = (std::f32::consts::PI * t as f32 / *t_max as f32).cos();
|
|
self.current_lr = eta_min + 0.5 * (self.initial_lr - eta_min) * (1.0 + cos_term);
|
|
}
|
|
LRSchedulerType::StepDecay { step_size, gamma } => {
|
|
// Reduce LR every step_size epochs
|
|
let decay_count = epoch / step_size;
|
|
self.current_lr = self.initial_lr * gamma.powi(decay_count as i32);
|
|
}
|
|
}
|
|
self.current_lr
|
|
}
|
|
|
|
/// Get current learning rate
|
|
#[must_use]
|
|
pub fn learning_rate(&self) -> f32 {
|
|
self.current_lr
|
|
}
|
|
}
|
|
|
|
/// Convert a slice of f32 values to bytes (little-endian).
|
|
fn f32_to_bytes(values: &[f32]) -> Vec<u8> {
|
|
values.iter().flat_map(|v| v.to_le_bytes()).collect()
|
|
}
|
|
|
|
/// Result type for training operations
|
|
pub type TrainingResult<T> = std::result::Result<T, TrainingError>;
|
|
|
|
/// Training-specific errors
|
|
#[derive(Debug, Clone)]
|
|
pub enum TrainingError {
|
|
/// Training was cancelled
|
|
Cancelled,
|
|
/// Data generation failed
|
|
DataGenerationFailed(String),
|
|
/// Model initialization failed
|
|
ModelInitFailed(String),
|
|
/// Training step failed
|
|
TrainingStepFailed(String),
|
|
/// Weight save failed
|
|
WeightSaveFailed(String),
|
|
/// Already training
|
|
AlreadyTraining,
|
|
/// Not initialized
|
|
NotInitialized,
|
|
}
|
|
|
|
impl std::fmt::Display for TrainingError {
|
|
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
|
match self {
|
|
Self::Cancelled => write!(f, "Training was cancelled"),
|
|
Self::DataGenerationFailed(msg) => write!(f, "Data generation failed: {msg}"),
|
|
Self::ModelInitFailed(msg) => write!(f, "Model initialization failed: {msg}"),
|
|
Self::TrainingStepFailed(msg) => write!(f, "Training step failed: {msg}"),
|
|
Self::WeightSaveFailed(msg) => write!(f, "Failed to save weights: {msg}"),
|
|
Self::AlreadyTraining => write!(f, "Training is already in progress"),
|
|
Self::NotInitialized => write!(f, "Trainer not initialized"),
|
|
}
|
|
}
|
|
}
|
|
|
|
impl std::error::Error for TrainingError {}
|
|
|
|
/// Trainer for FNO models
|
|
///
|
|
/// Manages the full training lifecycle including data generation,
|
|
/// model training, and progress reporting.
|
|
///
|
|
/// This trainer is generic over the compute backend, allowing training
|
|
/// on CPU, CUDA, Metal, or `ROCm` devices.
|
|
pub struct FnoTrainer<B: Backend<FloatElem = f32>> {
|
|
/// Training configuration
|
|
config: TrainingConfig,
|
|
/// PDE configuration
|
|
pde_config: PDEConfig,
|
|
/// Progress broadcast channel
|
|
progress_tx: broadcast::Sender<TrainingProgress>,
|
|
/// Flag indicating if training is active
|
|
is_training: Arc<AtomicBool>,
|
|
/// Flag to request cancellation
|
|
cancel_requested: Arc<AtomicBool>,
|
|
/// Current progress (for polling)
|
|
current_progress: std::sync::RwLock<TrainingProgress>,
|
|
/// Compute device
|
|
device: B::Device,
|
|
}
|
|
|
|
impl<B: Backend<FloatElem = f32>> FnoTrainer<B>
|
|
where
|
|
B::Device: Default,
|
|
{
|
|
/// Create a new FNO trainer with the default device for this backend
|
|
#[must_use]
|
|
pub fn new(config: TrainingConfig, pde_config: PDEConfig) -> Self {
|
|
Self::with_device(config, pde_config, B::Device::default())
|
|
}
|
|
}
|
|
|
|
impl<B: Backend<FloatElem = f32>> FnoTrainer<B> {
|
|
/// Get the device string for this backend type
|
|
fn backend_device_string() -> String {
|
|
// Use type_name to determine backend at runtime
|
|
let type_name = std::any::type_name::<B>();
|
|
|
|
if type_name.contains("CpuBackend") {
|
|
get_device_string_for_backend(BackendType::Cpu)
|
|
} else if type_name.contains("CudaBackend") {
|
|
get_device_string_for_backend(BackendType::Cuda)
|
|
} else if type_name.contains("MetalBackend") {
|
|
get_device_string_for_backend(BackendType::Metal)
|
|
} else if type_name.contains("RocmBackend") {
|
|
get_device_string_for_backend(BackendType::Rocm)
|
|
} else {
|
|
// Fallback to global detection
|
|
get_device_string()
|
|
}
|
|
}
|
|
|
|
/// Create a new FNO trainer with a specific device
|
|
#[must_use]
|
|
pub fn with_device(config: TrainingConfig, pde_config: PDEConfig, device: B::Device) -> Self {
|
|
let (progress_tx, _) = broadcast::channel(32);
|
|
let n_batches = config.n_train_samples.div_ceil(config.batch_size);
|
|
|
|
// Initialize progress with device info and learning rate
|
|
let mut initial_progress = TrainingProgress::new(config.epochs, n_batches);
|
|
initial_progress.device = Self::backend_device_string();
|
|
initial_progress.current_lr = config.learning_rate;
|
|
|
|
Self {
|
|
config,
|
|
pde_config,
|
|
progress_tx,
|
|
is_training: Arc::new(AtomicBool::new(false)),
|
|
cancel_requested: Arc::new(AtomicBool::new(false)),
|
|
current_progress: std::sync::RwLock::new(initial_progress),
|
|
device,
|
|
}
|
|
}
|
|
|
|
/// Check if training is currently active
|
|
#[must_use]
|
|
pub fn is_training(&self) -> bool {
|
|
self.is_training.load(Ordering::SeqCst)
|
|
}
|
|
|
|
/// Subscribe to progress updates
|
|
#[must_use]
|
|
pub fn subscribe(&self) -> broadcast::Receiver<TrainingProgress> {
|
|
self.progress_tx.subscribe()
|
|
}
|
|
|
|
/// Get current progress (for polling)
|
|
#[must_use]
|
|
pub fn get_progress(&self) -> TrainingProgress {
|
|
self.current_progress.read().unwrap().clone()
|
|
}
|
|
|
|
/// Request training cancellation
|
|
pub fn cancel(&self) {
|
|
self.cancel_requested.store(true, Ordering::SeqCst);
|
|
}
|
|
|
|
/// Run training synchronously (blocking)
|
|
///
|
|
/// Returns the path to saved weights on success.
|
|
pub fn train(&self) -> TrainingResult<String> {
|
|
if self
|
|
.is_training
|
|
.compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst)
|
|
.is_err()
|
|
{
|
|
return Err(TrainingError::AlreadyTraining);
|
|
}
|
|
|
|
// Reset cancellation flag
|
|
self.cancel_requested.store(false, Ordering::SeqCst);
|
|
|
|
let result = self.run_training_loop();
|
|
|
|
// Mark training as complete
|
|
self.is_training.store(false, Ordering::SeqCst);
|
|
|
|
result
|
|
}
|
|
|
|
/// Run the main training loop
|
|
fn run_training_loop(&self) -> TrainingResult<String> {
|
|
let start_time = Instant::now();
|
|
|
|
// Phase 1: Generate training data
|
|
self.update_progress(|p| {
|
|
p.status = TrainingStatus::GeneratingData {
|
|
samples_generated: 0,
|
|
total_samples: self.config.n_train_samples + self.config.n_val_samples,
|
|
};
|
|
});
|
|
|
|
let data_config =
|
|
DarcyDataConfig::default().with_resolution(self.pde_config.resolution as usize);
|
|
|
|
let generator = DarcyDataGenerator::with_config(data_config);
|
|
|
|
// Generate training samples
|
|
let train_samples =
|
|
self.generate_samples_with_progress(&generator, self.config.n_train_samples, 0)?;
|
|
|
|
// Generate validation samples
|
|
let val_samples = self.generate_samples_with_progress(
|
|
&generator,
|
|
self.config.n_val_samples,
|
|
self.config.n_train_samples,
|
|
)?;
|
|
|
|
if self.cancel_requested.load(Ordering::SeqCst) {
|
|
self.update_progress(|p| p.status = TrainingStatus::Cancelled);
|
|
return Err(TrainingError::Cancelled);
|
|
}
|
|
|
|
// Phase 2: Initialize model and optimizer
|
|
self.update_progress(|p| p.status = TrainingStatus::Training);
|
|
|
|
let pde_name = self
|
|
.pde_config
|
|
.pde_type
|
|
.name()
|
|
.to_lowercase()
|
|
.replace(' ', "_");
|
|
|
|
// Create initial model (using generic backend B)
|
|
let model = FNO2d::<B>::new_with_layers(
|
|
1, // in_channels (permeability field)
|
|
1, // out_channels (pressure field)
|
|
self.pde_config.model_width as usize,
|
|
self.pde_config.n_modes.0 as usize,
|
|
self.pde_config.n_layers as usize,
|
|
&self.device,
|
|
)
|
|
.map_err(|e| TrainingError::ModelInitFailed(e.to_string()))?;
|
|
|
|
// Extract initial weights and create optimizer
|
|
let mut weights = model.to_weights(&pde_name);
|
|
let mut params = extract_params(&weights);
|
|
let param_sizes: Vec<usize> = params.iter().map(std::vec::Vec::len).collect();
|
|
|
|
// Initialize AdamW optimizer with user-configured learning rate
|
|
let mut optimizer = AdamW::new(
|
|
¶m_sizes,
|
|
self.config.learning_rate,
|
|
0.9, // beta1
|
|
0.999, // beta2
|
|
0.01, // weight_decay
|
|
);
|
|
|
|
tracing::info!(
|
|
"Initialized AdamW optimizer: lr={}, params={}, total_elements={}",
|
|
optimizer.learning_rate(),
|
|
params.len(),
|
|
param_sizes.iter().sum::<usize>()
|
|
);
|
|
|
|
// Phase 3: Training loop with AdamW optimization
|
|
let n_batches = train_samples.len().div_ceil(self.config.batch_size);
|
|
let mut best_loss = f32::MAX;
|
|
let mut loss_history = Vec::with_capacity(self.config.epochs);
|
|
|
|
// Initialize early stopping (patience = 10 epochs, min_delta = 1e-4)
|
|
let mut early_stopping = EarlyStopping::new(10, 1e-4);
|
|
|
|
// Initialize LR scheduler (ReduceOnPlateau with factor=0.5, patience=5)
|
|
let mut lr_scheduler = LRScheduler::new(
|
|
LRSchedulerType::ReduceOnPlateau {
|
|
factor: 0.5,
|
|
patience: 5,
|
|
min_lr: 1e-6,
|
|
},
|
|
self.config.learning_rate,
|
|
);
|
|
|
|
// Epsilon for numerical gradient computation
|
|
const GRAD_EPS: f32 = 1e-4;
|
|
|
|
for epoch in 0..self.config.epochs {
|
|
if self.cancel_requested.load(Ordering::SeqCst) {
|
|
self.update_progress(|p| p.status = TrainingStatus::Cancelled);
|
|
return Err(TrainingError::Cancelled);
|
|
}
|
|
|
|
let epoch_start = Instant::now();
|
|
let mut epoch_loss = 0.0f32;
|
|
|
|
// Process batches
|
|
for batch_idx in 0..n_batches {
|
|
let batch_start = batch_idx * self.config.batch_size;
|
|
let batch_end = (batch_start + self.config.batch_size).min(train_samples.len());
|
|
let batch_samples = &train_samples[batch_start..batch_end];
|
|
|
|
// Create batch tensors
|
|
let batch = TrainingBatch::from_samples(batch_samples);
|
|
|
|
// Rebuild model with current parameters
|
|
update_weights(&mut weights, ¶ms);
|
|
let current_model = FNO2d::<B>::from_weights(&weights, &self.device)
|
|
.map_err(|e| TrainingError::TrainingStepFailed(e.to_string()))?;
|
|
|
|
// Compute current loss
|
|
let batch_loss = self.compute_batch_loss(¤t_model, &batch)?;
|
|
epoch_loss += batch_loss;
|
|
|
|
// Compute gradients using numerical differentiation
|
|
// For efficiency, we use stochastic parameter sampling
|
|
let grads = self.compute_numerical_gradients(
|
|
&mut weights,
|
|
¶ms,
|
|
&batch,
|
|
batch_loss,
|
|
GRAD_EPS,
|
|
)?;
|
|
|
|
// Apply AdamW optimizer step
|
|
optimizer.step(&mut params, &grads);
|
|
|
|
// Update progress
|
|
let elapsed = start_time.elapsed().as_secs_f32();
|
|
let samples_processed = epoch * train_samples.len() + batch_end;
|
|
let total_samples = self.config.epochs * train_samples.len();
|
|
let progress_frac = samples_processed as f32 / total_samples as f32;
|
|
let eta = if progress_frac > 0.01 {
|
|
elapsed / progress_frac - elapsed
|
|
} else {
|
|
0.0
|
|
};
|
|
|
|
// Capture current learning rate for progress reporting
|
|
let current_lr = optimizer.lr;
|
|
|
|
self.update_progress(|p| {
|
|
p.epoch = epoch;
|
|
p.batch = batch_idx;
|
|
p.total_batches = n_batches;
|
|
p.elapsed_seconds = elapsed;
|
|
p.eta_seconds = eta;
|
|
p.samples_per_sec = samples_processed as f32 / elapsed.max(0.01);
|
|
p.current_lr = current_lr;
|
|
});
|
|
}
|
|
|
|
// Compute average epoch loss
|
|
let avg_loss = epoch_loss / n_batches as f32;
|
|
loss_history.push(avg_loss);
|
|
|
|
if avg_loss < best_loss {
|
|
best_loss = avg_loss;
|
|
}
|
|
|
|
// Rebuild model for validation
|
|
update_weights(&mut weights, ¶ms);
|
|
let val_model = FNO2d::<B>::from_weights(&weights, &self.device)
|
|
.map_err(|e| TrainingError::TrainingStepFailed(e.to_string()))?;
|
|
|
|
// Validation loss
|
|
let val_loss = self.compute_validation_loss(&val_model, &val_samples)?;
|
|
|
|
// Update LR scheduler based on validation loss
|
|
let new_lr = lr_scheduler.step(epoch, Some(val_loss));
|
|
optimizer.lr = new_lr;
|
|
|
|
// Update progress with epoch results
|
|
let current_lr = new_lr;
|
|
let epochs_without_improvement = early_stopping.epochs_without_improvement();
|
|
self.update_progress(|p| {
|
|
p.loss = avg_loss;
|
|
p.best_loss = best_loss;
|
|
p.val_loss = Some(val_loss);
|
|
p.loss_history = loss_history.clone();
|
|
p.current_lr = current_lr;
|
|
});
|
|
|
|
// Log progress
|
|
let epoch_time = epoch_start.elapsed().as_secs_f32();
|
|
tracing::info!(
|
|
"Epoch {}/{}: loss={:.6}, val_loss={:.6}, time={:.2}s, lr={:.2e}, patience={}/10",
|
|
epoch + 1,
|
|
self.config.epochs,
|
|
avg_loss,
|
|
val_loss,
|
|
epoch_time,
|
|
new_lr,
|
|
epochs_without_improvement
|
|
);
|
|
|
|
// Check early stopping
|
|
if early_stopping.check(val_loss) {
|
|
tracing::info!(
|
|
"Early stopping triggered at epoch {} (no improvement for {} epochs)",
|
|
epoch + 1,
|
|
early_stopping.epochs_without_improvement()
|
|
);
|
|
break;
|
|
}
|
|
}
|
|
|
|
// Phase 4: Save final weights
|
|
update_weights(&mut weights, ¶ms);
|
|
let final_model = FNO2d::<B>::from_weights(&weights, &self.device)
|
|
.map_err(|e| TrainingError::TrainingStepFailed(e.to_string()))?;
|
|
let weights_path = self.save_weights(&final_model)?;
|
|
|
|
// Mark complete
|
|
self.update_progress(|p| {
|
|
p.status = TrainingStatus::Complete;
|
|
p.epoch = self.config.epochs;
|
|
});
|
|
|
|
Ok(weights_path)
|
|
}
|
|
|
|
/// Generate samples with progress tracking
|
|
fn generate_samples_with_progress(
|
|
&self,
|
|
generator: &DarcyDataGenerator,
|
|
n_samples: usize,
|
|
offset: usize,
|
|
) -> TrainingResult<Vec<DarcySample>> {
|
|
let total = self.config.n_train_samples + self.config.n_val_samples;
|
|
|
|
let samples = generator.generate_samples_with_progress(n_samples, |current, _| {
|
|
if current % 10 == 0 || current == n_samples {
|
|
self.update_progress(|p| {
|
|
p.status = TrainingStatus::GeneratingData {
|
|
samples_generated: offset + current,
|
|
total_samples: total,
|
|
};
|
|
});
|
|
}
|
|
});
|
|
|
|
Ok(samples)
|
|
}
|
|
|
|
/// Compute loss for a single batch
|
|
fn compute_batch_loss(&self, model: &FNO2d<B>, batch: &TrainingBatch) -> TrainingResult<f32> {
|
|
let batch_size = batch.batch_size;
|
|
let resolution = batch.resolution;
|
|
|
|
// Create input tensor [B, 1, H, W]
|
|
let input: GenericTensor<B, 4> = GenericTensor::from_slice(
|
|
&batch.inputs,
|
|
[batch_size, 1, resolution, resolution],
|
|
&self.device,
|
|
);
|
|
|
|
// Forward pass
|
|
let output = model.forward_4d(&input);
|
|
let output_data = output.to_vec();
|
|
|
|
// Compute MSE loss
|
|
let mut loss = 0.0f32;
|
|
for (pred, target) in output_data.iter().zip(batch.outputs.iter()) {
|
|
let diff = pred - target;
|
|
loss += diff * diff;
|
|
}
|
|
|
|
Ok(loss / batch.total_elements() as f32)
|
|
}
|
|
|
|
/// Compute validation loss
|
|
fn compute_validation_loss(
|
|
&self,
|
|
model: &FNO2d<B>,
|
|
samples: &[DarcySample],
|
|
) -> TrainingResult<f32> {
|
|
if samples.is_empty() {
|
|
return Ok(0.0);
|
|
}
|
|
|
|
let batch = TrainingBatch::from_samples(samples);
|
|
self.compute_batch_loss(model, &batch)
|
|
}
|
|
|
|
/// Compute numerical gradients using stochastic parameter perturbation.
|
|
///
|
|
/// For efficiency, we sample a subset of parameters to compute gradients
|
|
/// rather than computing gradients for all parameters every step.
|
|
/// This provides an unbiased gradient estimate while being computationally feasible.
|
|
///
|
|
/// # Arguments
|
|
/// * `weights` - Model weights structure (will be temporarily modified)
|
|
/// * `params` - Current parameter values
|
|
/// * `batch` - Training batch
|
|
/// * `base_loss` - Loss computed at current parameter values
|
|
/// * `eps` - Perturbation epsilon for finite differences
|
|
#[allow(clippy::too_many_arguments)]
|
|
fn compute_numerical_gradients(
|
|
&self,
|
|
weights: &mut FNO2dWeights,
|
|
params: &[Vec<f32>],
|
|
batch: &TrainingBatch,
|
|
base_loss: f32,
|
|
eps: f32,
|
|
) -> TrainingResult<Vec<Vec<f32>>> {
|
|
// Initialize gradients with zeros
|
|
let mut grads: Vec<Vec<f32>> = params.iter().map(|p| vec![0.0f32; p.len()]).collect();
|
|
|
|
// For efficiency, we use coordinate descent with stochastic sampling:
|
|
// Sample a subset of parameters per batch for gradient computation.
|
|
// This provides unbiased gradient estimates over time.
|
|
|
|
// Maximum number of parameter elements to sample per batch
|
|
const MAX_SAMPLES_PER_BATCH: usize = 256;
|
|
|
|
// Count total parameters
|
|
let total_params: usize = params.iter().map(std::vec::Vec::len).sum();
|
|
|
|
// Use a simple hash of the loss to seed randomness (deterministic within batch)
|
|
let seed = (base_loss.to_bits() ^ batch.batch_size as u32) as usize;
|
|
|
|
// Compute how many samples we can afford
|
|
let n_samples = MAX_SAMPLES_PER_BATCH.min(total_params);
|
|
|
|
// Sample parameter indices
|
|
let mut perturbed_params = params.to_vec();
|
|
|
|
for sample_idx in 0..n_samples {
|
|
// Deterministic "random" parameter selection using simple hashing
|
|
let flat_idx = (seed.wrapping_mul(31).wrapping_add(sample_idx * 17)) % total_params;
|
|
|
|
// Convert flat index to (param_group, param_idx)
|
|
let mut remaining = flat_idx;
|
|
let mut param_group = 0;
|
|
while param_group < params.len() && remaining >= params[param_group].len() {
|
|
remaining -= params[param_group].len();
|
|
param_group += 1;
|
|
}
|
|
|
|
if param_group >= params.len() {
|
|
continue;
|
|
}
|
|
|
|
let param_idx = remaining;
|
|
|
|
// Compute gradient using central differences: (f(x+eps) - f(x-eps)) / (2*eps)
|
|
// This is more accurate than forward differences
|
|
|
|
// Perturb parameter positively
|
|
let original_value = perturbed_params[param_group][param_idx];
|
|
perturbed_params[param_group][param_idx] = original_value + eps;
|
|
update_weights(weights, &perturbed_params);
|
|
|
|
let perturbed_model = FNO2d::<B>::from_weights(weights, &self.device)
|
|
.map_err(|e| TrainingError::TrainingStepFailed(e.to_string()))?;
|
|
let loss_plus = self.compute_batch_loss(&perturbed_model, batch)?;
|
|
|
|
// Perturb parameter negatively
|
|
perturbed_params[param_group][param_idx] = original_value - eps;
|
|
update_weights(weights, &perturbed_params);
|
|
|
|
let perturbed_model = FNO2d::<B>::from_weights(weights, &self.device)
|
|
.map_err(|e| TrainingError::TrainingStepFailed(e.to_string()))?;
|
|
let loss_minus = self.compute_batch_loss(&perturbed_model, batch)?;
|
|
|
|
// Restore original value
|
|
perturbed_params[param_group][param_idx] = original_value;
|
|
|
|
// Compute gradient using central differences
|
|
let grad = (loss_plus - loss_minus) / (2.0 * eps);
|
|
|
|
// Scale gradient by ratio of total to sampled parameters for unbiased estimate
|
|
let scale = total_params as f32 / n_samples as f32;
|
|
grads[param_group][param_idx] = grad * scale;
|
|
}
|
|
|
|
// Restore weights to original state
|
|
update_weights(weights, params);
|
|
|
|
Ok(grads)
|
|
}
|
|
|
|
/// Save trained weights to `SafeTensors` format
|
|
fn save_weights(&self, model: &FNO2d<B>) -> TrainingResult<String> {
|
|
use rtx_hub::safetensors::{SafeTensorsBuilder, SafeTensorsDType};
|
|
|
|
// Create weights directory
|
|
let weights_dir = dirs::data_local_dir()
|
|
.unwrap_or_else(std::env::temp_dir)
|
|
.join("rustytorch/neural-operator");
|
|
std::fs::create_dir_all(&weights_dir)
|
|
.map_err(|e| TrainingError::WeightSaveFailed(e.to_string()))?;
|
|
|
|
let pde_name = self
|
|
.pde_config
|
|
.pde_type
|
|
.name()
|
|
.to_lowercase()
|
|
.replace(' ', "_");
|
|
let weights_path = weights_dir.join(format!(
|
|
"fno_{}_{}.safetensors",
|
|
pde_name, self.pde_config.resolution
|
|
));
|
|
|
|
// Extract weights from model
|
|
let weights = model.to_weights(&pde_name);
|
|
|
|
// Build SafeTensors file
|
|
let builder = SafeTensorsBuilder::new()
|
|
// Metadata
|
|
.with_metadata("model_type", "FNO2d")
|
|
.with_metadata("pde_type", &weights.config.pde_type)
|
|
.with_metadata("resolution", self.pde_config.resolution.to_string())
|
|
.with_metadata("in_channels", weights.config.in_channels.to_string())
|
|
.with_metadata("out_channels", weights.config.out_channels.to_string())
|
|
.with_metadata("width", weights.config.width.to_string())
|
|
.with_metadata("n_layers", weights.config.n_layers.to_string())
|
|
.with_metadata("n_modes_h", weights.config.n_modes.0.to_string())
|
|
.with_metadata("n_modes_w", weights.config.n_modes.1.to_string())
|
|
// Lifting MLP weights
|
|
.add_tensor(
|
|
"lifting.fcs.0.weight",
|
|
SafeTensorsDType::F32,
|
|
vec![weights.config.width * 2, weights.config.in_channels + 2],
|
|
f32_to_bytes(&weights.lifting_fc1_weight),
|
|
)
|
|
.add_tensor(
|
|
"lifting.fcs.0.bias",
|
|
SafeTensorsDType::F32,
|
|
vec![weights.config.width * 2],
|
|
f32_to_bytes(&weights.lifting_fc1_bias),
|
|
)
|
|
.add_tensor(
|
|
"lifting.fcs.1.weight",
|
|
SafeTensorsDType::F32,
|
|
vec![weights.config.width, weights.config.width * 2],
|
|
f32_to_bytes(&weights.lifting_fc2_weight),
|
|
)
|
|
.add_tensor(
|
|
"lifting.fcs.1.bias",
|
|
SafeTensorsDType::F32,
|
|
vec![weights.config.width],
|
|
f32_to_bytes(&weights.lifting_fc2_bias),
|
|
);
|
|
|
|
// Add spectral conv weights for each layer
|
|
let mut builder = builder;
|
|
for (i, sw) in weights.spectral_weights.iter().enumerate() {
|
|
let shape = vec![
|
|
weights.config.width,
|
|
weights.config.width,
|
|
weights.config.n_modes.0,
|
|
weights.config.n_modes.1,
|
|
];
|
|
builder = builder
|
|
.add_tensor(
|
|
format!("spectral_conv.{i}.weights1_real"),
|
|
SafeTensorsDType::F32,
|
|
shape.clone(),
|
|
f32_to_bytes(&sw.weights1_real),
|
|
)
|
|
.add_tensor(
|
|
format!("spectral_conv.{i}.weights1_imag"),
|
|
SafeTensorsDType::F32,
|
|
shape.clone(),
|
|
f32_to_bytes(&sw.weights1_imag),
|
|
)
|
|
.add_tensor(
|
|
format!("spectral_conv.{i}.weights2_real"),
|
|
SafeTensorsDType::F32,
|
|
shape.clone(),
|
|
f32_to_bytes(&sw.weights2_real),
|
|
)
|
|
.add_tensor(
|
|
format!("spectral_conv.{i}.weights2_imag"),
|
|
SafeTensorsDType::F32,
|
|
shape,
|
|
f32_to_bytes(&sw.weights2_imag),
|
|
);
|
|
}
|
|
|
|
// Add 1x1 conv weights for each layer
|
|
for (i, (weight, bias)) in weights.conv_weights.iter().enumerate() {
|
|
builder = builder
|
|
.add_tensor(
|
|
format!("conv.{i}.weight"),
|
|
SafeTensorsDType::F32,
|
|
vec![weights.config.width, weights.config.width],
|
|
f32_to_bytes(weight),
|
|
)
|
|
.add_tensor(
|
|
format!("conv.{i}.bias"),
|
|
SafeTensorsDType::F32,
|
|
vec![weights.config.width],
|
|
f32_to_bytes(bias),
|
|
);
|
|
}
|
|
|
|
// Add projection layer weights
|
|
const PROJECTION_HIDDEN: usize = 128;
|
|
let (proj0_weight, proj0_bias) = &weights.projection_weights[0];
|
|
let (proj1_weight, proj1_bias) = &weights.projection_weights[1];
|
|
|
|
builder = builder
|
|
.add_tensor(
|
|
"projection.0.weight",
|
|
SafeTensorsDType::F32,
|
|
vec![PROJECTION_HIDDEN, weights.config.width],
|
|
f32_to_bytes(proj0_weight),
|
|
)
|
|
.add_tensor(
|
|
"projection.0.bias",
|
|
SafeTensorsDType::F32,
|
|
vec![PROJECTION_HIDDEN],
|
|
f32_to_bytes(proj0_bias),
|
|
)
|
|
.add_tensor(
|
|
"projection.1.weight",
|
|
SafeTensorsDType::F32,
|
|
vec![weights.config.out_channels, PROJECTION_HIDDEN],
|
|
f32_to_bytes(proj1_weight),
|
|
)
|
|
.add_tensor(
|
|
"projection.1.bias",
|
|
SafeTensorsDType::F32,
|
|
vec![weights.config.out_channels],
|
|
f32_to_bytes(proj1_bias),
|
|
);
|
|
|
|
// Build and write file
|
|
let file_bytes = builder
|
|
.build()
|
|
.map_err(|e| TrainingError::WeightSaveFailed(e.to_string()))?;
|
|
|
|
std::fs::write(&weights_path, file_bytes)
|
|
.map_err(|e| TrainingError::WeightSaveFailed(e.to_string()))?;
|
|
|
|
tracing::info!("Weights saved to {:?}", weights_path);
|
|
|
|
Ok(weights_path.to_string_lossy().to_string())
|
|
}
|
|
|
|
/// Update progress and broadcast
|
|
fn update_progress<F>(&self, updater: F)
|
|
where
|
|
F: FnOnce(&mut TrainingProgress),
|
|
{
|
|
let mut progress = self.current_progress.write().unwrap();
|
|
updater(&mut progress);
|
|
|
|
// Broadcast update (ignore send errors - no receivers)
|
|
let _ = self.progress_tx.send(progress.clone());
|
|
}
|
|
}
|
|
|
|
// ============================================================================
|
|
// Dynamic Backend Dispatch
|
|
// ============================================================================
|
|
|
|
/// Trait object for dynamic dispatch of training operations.
|
|
///
|
|
/// This enables runtime selection of compute backend without requiring
|
|
/// the caller to know the specific backend type.
|
|
pub trait DynamicTrainer: Send + Sync {
|
|
/// Start training and return the weights path on success
|
|
fn train(&self) -> TrainingResult<String>;
|
|
/// Get current training progress
|
|
fn get_progress(&self) -> TrainingProgress;
|
|
/// Cancel training
|
|
fn cancel(&self);
|
|
/// Check if training is active
|
|
fn is_training(&self) -> bool;
|
|
}
|
|
|
|
impl<B: Backend<FloatElem = f32> + Send + Sync + 'static> DynamicTrainer for FnoTrainer<B>
|
|
where
|
|
B::Device: Send + Sync,
|
|
{
|
|
fn train(&self) -> TrainingResult<String> {
|
|
FnoTrainer::train(self)
|
|
}
|
|
|
|
fn get_progress(&self) -> TrainingProgress {
|
|
FnoTrainer::get_progress(self)
|
|
}
|
|
|
|
fn cancel(&self) {
|
|
FnoTrainer::cancel(self);
|
|
}
|
|
|
|
fn is_training(&self) -> bool {
|
|
FnoTrainer::is_training(self)
|
|
}
|
|
}
|
|
|
|
/// Training session handle for async operations.
|
|
///
|
|
/// This struct uses runtime backend selection to automatically choose
|
|
/// the best available compute device (CUDA > Metal > `ROCm` > CPU).
|
|
pub struct TrainingSession {
|
|
trainer: Arc<dyn DynamicTrainer>,
|
|
}
|
|
|
|
impl TrainingSession {
|
|
/// Create a new training session with automatic backend selection.
|
|
///
|
|
/// The backend is selected based on availability:
|
|
/// 1. CUDA (if nvidia-smi available and feature enabled)
|
|
/// 2. Metal (on macOS, if feature enabled)
|
|
/// 3. `ROCm` (if rocminfo available and feature enabled)
|
|
/// 4. CPU (always available)
|
|
pub fn new(config: TrainingConfig, pde_config: PDEConfig) -> Self {
|
|
let backend_type = select_best_backend();
|
|
tracing::info!("Training session using backend: {}", backend_type);
|
|
|
|
let trainer: Arc<dyn DynamicTrainer> = match backend_type {
|
|
#[cfg(feature = "cuda")]
|
|
BackendType::Cuda => {
|
|
tracing::info!("Initializing CUDA backend for training");
|
|
Arc::new(FnoTrainer::<CudaBackend>::new(config, pde_config))
|
|
}
|
|
#[cfg(feature = "metal")]
|
|
BackendType::Metal => {
|
|
tracing::info!("Initializing Metal backend for training");
|
|
Arc::new(FnoTrainer::<MetalBackend>::new(config, pde_config))
|
|
}
|
|
#[cfg(feature = "rocm")]
|
|
BackendType::Rocm => {
|
|
tracing::info!("Initializing ROCm backend for training");
|
|
Arc::new(FnoTrainer::<RocmBackend>::new(config, pde_config))
|
|
}
|
|
_ => {
|
|
tracing::info!("Initializing CPU backend for training");
|
|
Arc::new(FnoTrainer::<CpuBackend>::new(config, pde_config))
|
|
}
|
|
};
|
|
|
|
Self { trainer }
|
|
}
|
|
|
|
/// Create a training session with a specific backend.
|
|
///
|
|
/// This is useful for testing or when you want to force a specific backend.
|
|
#[must_use]
|
|
pub fn with_backend<B: Backend<FloatElem = f32> + Send + Sync + 'static>(
|
|
config: TrainingConfig,
|
|
pde_config: PDEConfig,
|
|
) -> Self
|
|
where
|
|
B::Device: Default + Send + Sync,
|
|
{
|
|
Self {
|
|
trainer: Arc::new(FnoTrainer::<B>::new(config, pde_config)),
|
|
}
|
|
}
|
|
|
|
/// Start training in background
|
|
#[must_use]
|
|
pub fn start(&self) -> tokio::task::JoinHandle<TrainingResult<String>> {
|
|
let trainer = Arc::clone(&self.trainer);
|
|
tokio::task::spawn_blocking(move || trainer.train())
|
|
}
|
|
|
|
/// Get current progress
|
|
#[must_use]
|
|
pub fn get_progress(&self) -> TrainingProgress {
|
|
self.trainer.get_progress()
|
|
}
|
|
|
|
/// Cancel training
|
|
pub fn cancel(&self) {
|
|
self.trainer.cancel();
|
|
}
|
|
|
|
/// Check if training is active
|
|
#[must_use]
|
|
pub fn is_training(&self) -> bool {
|
|
self.trainer.is_training()
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_trainer_creation() {
|
|
let config = TrainingConfig::quick();
|
|
let pde_config = PDEConfig::darcy(32);
|
|
// Explicitly use CpuBackend for tests
|
|
let trainer = FnoTrainer::<CpuBackend>::new(config, pde_config);
|
|
|
|
assert!(!trainer.is_training());
|
|
}
|
|
|
|
#[test]
|
|
fn test_progress_update() {
|
|
let config = TrainingConfig::quick();
|
|
let pde_config = PDEConfig::darcy(32);
|
|
let trainer = FnoTrainer::<CpuBackend>::new(config, pde_config);
|
|
|
|
let initial_progress = trainer.get_progress();
|
|
assert_eq!(initial_progress.epoch, 0);
|
|
assert_eq!(initial_progress.status, TrainingStatus::NotStarted);
|
|
|
|
trainer.update_progress(|p| {
|
|
p.epoch = 5;
|
|
p.loss = 0.01;
|
|
});
|
|
|
|
let updated_progress = trainer.get_progress();
|
|
assert_eq!(updated_progress.epoch, 5);
|
|
assert!((updated_progress.loss - 0.01).abs() < 1e-6);
|
|
}
|
|
|
|
#[test]
|
|
fn test_cancel_flag() {
|
|
let config = TrainingConfig::quick();
|
|
let pde_config = PDEConfig::darcy(32);
|
|
let trainer = FnoTrainer::<CpuBackend>::new(config, pde_config);
|
|
|
|
assert!(!trainer.cancel_requested.load(Ordering::SeqCst));
|
|
trainer.cancel();
|
|
assert!(trainer.cancel_requested.load(Ordering::SeqCst));
|
|
}
|
|
|
|
// ============================================================================
|
|
// TDD Tests for AdamW Optimizer (RED-GREEN-REFACTOR)
|
|
// ============================================================================
|
|
|
|
#[test]
|
|
fn test_adamw_optimizer_creation() {
|
|
// RED: Test that AdamW optimizer can be created with config
|
|
let param_sizes = vec![4, 2];
|
|
let learning_rate = 0.001;
|
|
let beta1 = 0.9;
|
|
let beta2 = 0.999;
|
|
let weight_decay = 0.01;
|
|
|
|
let optimizer = AdamW::new(¶m_sizes, learning_rate, beta1, beta2, weight_decay);
|
|
|
|
// GREEN: Verify optimizer was created successfully
|
|
assert_eq!(optimizer.lr, learning_rate);
|
|
assert_eq!(optimizer.beta1, beta1);
|
|
assert_eq!(optimizer.beta2, beta2);
|
|
assert_eq!(optimizer.weight_decay, weight_decay);
|
|
assert_eq!(optimizer.eps, 1e-8);
|
|
assert_eq!(optimizer.t, 0, "Initial timestep should be 0");
|
|
|
|
// Verify momentum and variance are initialized to zeros
|
|
assert_eq!(optimizer.m.len(), param_sizes.len());
|
|
assert_eq!(optimizer.v.len(), param_sizes.len());
|
|
for (i, &size) in param_sizes.iter().enumerate() {
|
|
assert_eq!(optimizer.m[i].len(), size);
|
|
assert_eq!(optimizer.v[i].len(), size);
|
|
assert!(optimizer.m[i].iter().all(|&x| x == 0.0));
|
|
assert!(optimizer.v[i].iter().all(|&x| x == 0.0));
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_adamw_optimizer_step_updates_weights() {
|
|
// RED: Test that optimizer step updates weights
|
|
let param_sizes = vec![4, 2];
|
|
let mut optimizer = AdamW::new(¶m_sizes, 0.1, 0.9, 0.999, 0.01);
|
|
|
|
let mut params = vec![vec![1.0, 2.0, 3.0, 4.0], vec![5.0, 6.0]];
|
|
|
|
let grads = vec![vec![0.1, 0.2, 0.3, 0.4], vec![0.5, 0.6]];
|
|
|
|
let original_params = params.clone();
|
|
|
|
// GREEN: Perform optimization step
|
|
optimizer.step(&mut params, &grads);
|
|
|
|
// All parameters should have changed
|
|
for (i, (new_param, old_param)) in params.iter().zip(original_params.iter()).enumerate() {
|
|
for (j, (new_val, old_val)) in new_param.iter().zip(old_param.iter()).enumerate() {
|
|
assert!(
|
|
(new_val - old_val).abs() > 1e-8,
|
|
"Parameter [{i}][{j}] should have changed: {old_val} -> {new_val}"
|
|
);
|
|
}
|
|
}
|
|
|
|
// All parameters should have decreased (positive gradient with descent)
|
|
for (i, (new_param, old_param)) in params.iter().zip(original_params.iter()).enumerate() {
|
|
for (j, (new_val, old_val)) in new_param.iter().zip(old_param.iter()).enumerate() {
|
|
assert!(
|
|
*new_val < *old_val,
|
|
"Parameter [{i}][{j}] should decrease: {old_val} -> {new_val}"
|
|
);
|
|
}
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_adamw_maintains_momentum_state() {
|
|
// RED: Test that momentum state is maintained across batches
|
|
let param_sizes = vec![4];
|
|
let mut optimizer = AdamW::new(¶m_sizes, 0.01, 0.9, 0.999, 0.01);
|
|
|
|
let mut params = vec![vec![1.0, 1.0, 1.0, 1.0]];
|
|
let grads = vec![vec![0.1, 0.1, 0.1, 0.1]];
|
|
|
|
// GREEN: First step
|
|
optimizer.step(&mut params, &grads);
|
|
|
|
// After first step, momentum should be non-zero
|
|
assert!(
|
|
optimizer.m[0].iter().any(|&x| x.abs() > 1e-8),
|
|
"Momentum should be non-zero after first step"
|
|
);
|
|
|
|
// Expected momentum after first step: m = β₁ * 0 + (1 - β₁) * 0.1 = 0.1 * 0.1 = 0.01
|
|
let expected_m_first = (1.0 - 0.9) * 0.1;
|
|
assert!(
|
|
(optimizer.m[0][0] - expected_m_first).abs() < 1e-6,
|
|
"Momentum after first step should be {}, got {}",
|
|
expected_m_first,
|
|
optimizer.m[0][0]
|
|
);
|
|
|
|
// Second step with same gradient
|
|
let m_before_second = optimizer.m[0][0];
|
|
optimizer.step(&mut params, &grads);
|
|
|
|
// Momentum should have accumulated: m = β₁ * m_prev + (1 - β₁) * g
|
|
let expected_m_second = 0.9 * m_before_second + (1.0 - 0.9) * 0.1;
|
|
assert!(
|
|
(optimizer.m[0][0] - expected_m_second).abs() < 1e-6,
|
|
"Momentum after second step should be {}, got {}",
|
|
expected_m_second,
|
|
optimizer.m[0][0]
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_adamw_maintains_variance_state() {
|
|
// RED: Test that variance state is maintained across batches
|
|
let param_sizes = vec![4];
|
|
let mut optimizer = AdamW::new(¶m_sizes, 0.01, 0.9, 0.999, 0.01);
|
|
|
|
let mut params = vec![vec![1.0, 1.0, 1.0, 1.0]];
|
|
let grads = vec![vec![0.1, 0.1, 0.1, 0.1]];
|
|
|
|
// GREEN: First step
|
|
optimizer.step(&mut params, &grads);
|
|
|
|
// After first step, variance should be non-zero
|
|
assert!(
|
|
optimizer.v[0].iter().any(|&x| x.abs() > 1e-8),
|
|
"Variance should be non-zero after first step"
|
|
);
|
|
|
|
// Expected variance after first step: v = β₂ * 0 + (1 - β₂) * g² = 0.001 * 0.01 = 0.00001
|
|
let expected_v_first = (1.0 - 0.999) * 0.1 * 0.1;
|
|
assert!(
|
|
(optimizer.v[0][0] - expected_v_first).abs() < 1e-8,
|
|
"Variance after first step should be {}, got {}",
|
|
expected_v_first,
|
|
optimizer.v[0][0]
|
|
);
|
|
|
|
// Second step
|
|
let v_before_second = optimizer.v[0][0];
|
|
optimizer.step(&mut params, &grads);
|
|
|
|
// Variance should have accumulated
|
|
let expected_v_second = 0.999 * v_before_second + (1.0 - 0.999) * 0.1 * 0.1;
|
|
assert!(
|
|
(optimizer.v[0][0] - expected_v_second).abs() < 1e-8,
|
|
"Variance after second step should be {}, got {}",
|
|
expected_v_second,
|
|
optimizer.v[0][0]
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_adamw_timestep_increments() {
|
|
// RED: Test that timestep counter increments with each step
|
|
let param_sizes = vec![2];
|
|
let mut optimizer = AdamW::new(¶m_sizes, 0.01, 0.9, 0.999, 0.01);
|
|
let mut params = vec![vec![1.0, 2.0]];
|
|
let grads = vec![vec![0.1, 0.1]];
|
|
|
|
// GREEN: Verify timestep increments
|
|
assert_eq!(optimizer.t, 0, "Initial timestep should be 0");
|
|
|
|
optimizer.step(&mut params, &grads);
|
|
assert_eq!(optimizer.t, 1, "Timestep should be 1 after first step");
|
|
|
|
optimizer.step(&mut params, &grads);
|
|
assert_eq!(optimizer.t, 2, "Timestep should be 2 after second step");
|
|
|
|
optimizer.step(&mut params, &grads);
|
|
assert_eq!(optimizer.t, 3, "Timestep should be 3 after third step");
|
|
}
|
|
|
|
#[test]
|
|
fn test_adamw_loss_decreases_faster_than_sgd() {
|
|
// RED: Test that loss decreases faster with AdamW than with manual SGD
|
|
// This demonstrates the adaptive learning rate benefit
|
|
|
|
// Setup: Two identical training scenarios
|
|
let param_sizes = vec![10]; // Single weight vector
|
|
let mut params_adamw = vec![vec![1.0; 10]];
|
|
let mut params_sgd = vec![vec![1.0; 10]];
|
|
|
|
let learning_rate = 0.01;
|
|
let mut optimizer_adamw = AdamW::new(¶m_sizes, learning_rate, 0.9, 0.999, 0.0);
|
|
|
|
// Simulate 20 steps with constant gradient
|
|
let grads = vec![vec![0.1; 10]];
|
|
|
|
// GREEN: Run optimization
|
|
for _ in 0..20 {
|
|
// AdamW step
|
|
optimizer_adamw.step(&mut params_adamw, &grads);
|
|
|
|
// Manual SGD step: param -= lr * grad
|
|
for (param_vec, grad_vec) in params_sgd.iter_mut().zip(grads.iter()) {
|
|
for (p, &g) in param_vec.iter_mut().zip(grad_vec.iter()) {
|
|
*p -= learning_rate * g;
|
|
}
|
|
}
|
|
}
|
|
|
|
// Compute "loss" as squared distance from origin (simulating convergence to zero)
|
|
let loss_adamw: f32 = params_adamw[0].iter().map(|&x| x * x).sum();
|
|
let loss_sgd: f32 = params_sgd[0].iter().map(|&x| x * x).sum();
|
|
|
|
// AdamW should converge faster due to adaptive learning and momentum
|
|
assert!(
|
|
loss_adamw < loss_sgd,
|
|
"AdamW loss ({}) should be lower than SGD loss ({}) after 20 steps",
|
|
loss_adamw,
|
|
loss_sgd
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_adamw_bias_correction() {
|
|
// RED: Test that bias correction is applied properly in early steps
|
|
let param_sizes = vec![2];
|
|
let mut optimizer = AdamW::new(¶m_sizes, 1.0, 0.9, 0.999, 0.0);
|
|
let mut params = vec![vec![1.0, 1.0]];
|
|
let grads = vec![vec![1.0, 1.0]];
|
|
|
|
let params_before = params.clone();
|
|
|
|
// GREEN: First step
|
|
optimizer.step(&mut params, &grads);
|
|
|
|
// Without bias correction, the update would be very small
|
|
// With bias correction: m_hat = m / (1 - β₁^t) = 0.1 / 0.1 = 1.0
|
|
// The parameter change should be significant due to bias correction
|
|
let change = (params[0][0] - params_before[0][0]).abs();
|
|
assert!(
|
|
change > 0.5,
|
|
"First step should have large update due to bias correction, change = {}",
|
|
change
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_adamw_weight_decay_applied() {
|
|
// RED: Test that weight decay is applied correctly
|
|
let param_sizes = vec![4];
|
|
let weight_decay = 0.1; // Significant weight decay for testing
|
|
let mut optimizer = AdamW::new(¶m_sizes, 0.001, 0.9, 0.999, weight_decay);
|
|
|
|
let mut params = vec![vec![1.0, 2.0, 3.0, 4.0]];
|
|
let grads = vec![vec![0.0, 0.0, 0.0, 0.0]]; // Zero gradient
|
|
|
|
let params_before = params.clone();
|
|
|
|
// GREEN: With zero gradient and weight decay, parameters should shrink toward zero
|
|
optimizer.step(&mut params, &grads);
|
|
|
|
for (i, (&new_val, &old_val)) in params[0].iter().zip(params_before[0].iter()).enumerate() {
|
|
assert!(
|
|
new_val < old_val,
|
|
"Parameter [{}] should decrease due to weight decay: {} -> {}",
|
|
i,
|
|
old_val,
|
|
new_val
|
|
);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_adamw_with_default_hyperparameters() {
|
|
// RED: Test that default hyperparameters produce reasonable behavior
|
|
let param_sizes = vec![4, 2];
|
|
let optimizer = AdamW::with_defaults(¶m_sizes);
|
|
|
|
// GREEN: Verify defaults
|
|
assert_eq!(optimizer.lr, 1e-3);
|
|
assert_eq!(optimizer.beta1, 0.9);
|
|
assert_eq!(optimizer.beta2, 0.999);
|
|
assert_eq!(optimizer.weight_decay, 0.01);
|
|
}
|
|
|
|
#[test]
|
|
fn test_adamw_integration_with_training_config() {
|
|
// RED: Test that AdamW integrates with training configuration
|
|
let config = TrainingConfig {
|
|
epochs: 2,
|
|
batch_size: 4,
|
|
learning_rate: 0.002, // Custom learning rate
|
|
n_train_samples: 8,
|
|
n_val_samples: 4,
|
|
};
|
|
|
|
let param_sizes = vec![10, 5];
|
|
let optimizer = AdamW::new(¶m_sizes, config.learning_rate, 0.9, 0.999, 0.01);
|
|
|
|
// GREEN: Verify optimizer uses config learning rate
|
|
assert_eq!(optimizer.lr, config.learning_rate);
|
|
}
|
|
|
|
#[test]
|
|
fn test_extract_update_weights_roundtrip() {
|
|
// Test that extract_params and update_weights are inverses
|
|
let device = CpuDevice::default();
|
|
let model = FNO2d::<CpuBackend>::new_with_layers(1, 1, 8, 4, 2, &device)
|
|
.expect("Failed to create model");
|
|
|
|
let mut weights = model.to_weights("test");
|
|
let original_first_weight = weights.lifting_fc1_weight[0];
|
|
|
|
let params = extract_params(&weights);
|
|
|
|
// Verify extracted parameters match
|
|
assert!(
|
|
(params[0][0] - original_first_weight).abs() < 1e-6,
|
|
"Extracted params should match original"
|
|
);
|
|
|
|
// Modify a parameter
|
|
let mut modified_params = params.clone();
|
|
modified_params[0][0] += 1.0;
|
|
|
|
// Update weights and check the modification took effect
|
|
update_weights(&mut weights, &modified_params);
|
|
|
|
assert!(
|
|
(weights.lifting_fc1_weight[0] - original_first_weight - 1.0).abs() < 1e-6,
|
|
"Weight modification should persist"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_early_stopping() {
|
|
let mut early_stop = EarlyStopping::new(3, 1e-4);
|
|
|
|
// First check with good loss - should not stop
|
|
assert!(!early_stop.check(1.0));
|
|
assert_eq!(early_stop.epochs_without_improvement(), 0);
|
|
|
|
// Improvement - counter stays at 0
|
|
assert!(!early_stop.check(0.5));
|
|
assert_eq!(early_stop.epochs_without_improvement(), 0);
|
|
|
|
// No improvement - counter increases
|
|
assert!(!early_stop.check(0.6));
|
|
assert_eq!(early_stop.epochs_without_improvement(), 1);
|
|
|
|
assert!(!early_stop.check(0.7));
|
|
assert_eq!(early_stop.epochs_without_improvement(), 2);
|
|
|
|
// Third no improvement - triggers early stop
|
|
assert!(early_stop.check(0.8));
|
|
|
|
// Improvement resets counter
|
|
let mut early_stop2 = EarlyStopping::new(3, 1e-4);
|
|
early_stop2.check(1.0);
|
|
early_stop2.check(1.1); // no improvement
|
|
early_stop2.check(1.2); // no improvement
|
|
assert!(!early_stop2.check(0.5)); // improvement - resets
|
|
assert_eq!(early_stop2.epochs_without_improvement(), 0);
|
|
}
|
|
|
|
#[test]
|
|
fn test_lr_scheduler_reduce_on_plateau() {
|
|
let mut scheduler = LRScheduler::new(
|
|
LRSchedulerType::ReduceOnPlateau {
|
|
factor: 0.5,
|
|
patience: 2,
|
|
min_lr: 1e-6,
|
|
},
|
|
0.01,
|
|
);
|
|
|
|
// Initial LR
|
|
assert!((scheduler.learning_rate() - 0.01).abs() < 1e-8);
|
|
|
|
// Good loss - no change
|
|
scheduler.step(0, Some(1.0));
|
|
assert!((scheduler.learning_rate() - 0.01).abs() < 1e-8);
|
|
|
|
// No improvement for 2 epochs - triggers reduction
|
|
scheduler.step(1, Some(1.1)); // worse
|
|
scheduler.step(2, Some(1.2)); // worse again, triggers reduction
|
|
assert!((scheduler.learning_rate() - 0.005).abs() < 1e-8);
|
|
}
|
|
|
|
#[test]
|
|
fn test_lr_scheduler_cosine_annealing() {
|
|
let mut scheduler = LRScheduler::new(
|
|
LRSchedulerType::CosineAnnealing {
|
|
t_max: 100,
|
|
eta_min: 0.0001,
|
|
},
|
|
0.01,
|
|
);
|
|
|
|
// At epoch 0, LR should be at maximum (cos(0) = 1)
|
|
// lr = eta_min + 0.5 * (initial_lr - eta_min) * (1 + 1) = initial_lr
|
|
scheduler.step(0, None);
|
|
assert!((scheduler.learning_rate() - 0.01).abs() < 1e-6);
|
|
|
|
// At epoch t_max/2 (50), LR should be at midpoint (cos(π/2) = 0)
|
|
// lr = eta_min + 0.5 * (initial_lr - eta_min) * (1 + 0) = (initial_lr + eta_min) / 2
|
|
scheduler.step(50, None);
|
|
let expected_mid = (0.01 + 0.0001) / 2.0; // 0.00505
|
|
assert!(
|
|
(scheduler.learning_rate() - expected_mid).abs() < 1e-6,
|
|
"At midpoint: expected {}, got {}",
|
|
expected_mid,
|
|
scheduler.learning_rate()
|
|
);
|
|
|
|
// At epoch 99 (just before restart), LR should be near minimum
|
|
scheduler.step(99, None);
|
|
// t=99, cos(π*99/100) ≈ cos(0.99π) ≈ -0.9998
|
|
// lr ≈ eta_min + 0.5 * 0.0099 * (1 - 0.9998) ≈ 0.0001
|
|
assert!(
|
|
scheduler.learning_rate() < 0.001,
|
|
"Near end of cycle: expected near eta_min, got {}",
|
|
scheduler.learning_rate()
|
|
);
|
|
|
|
// At epoch t_max (100), cycle restarts so LR goes back to max
|
|
// Because epoch % t_max = 0
|
|
scheduler.step(100, None);
|
|
assert!(
|
|
(scheduler.learning_rate() - 0.01).abs() < 1e-6,
|
|
"Cycle restart: expected {}, got {}",
|
|
0.01,
|
|
scheduler.learning_rate()
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_lr_scheduler_step_decay() {
|
|
let mut scheduler = LRScheduler::new(
|
|
LRSchedulerType::StepDecay {
|
|
step_size: 10,
|
|
gamma: 0.1,
|
|
},
|
|
0.01,
|
|
);
|
|
|
|
// Epochs 0-9: LR = 0.01
|
|
scheduler.step(0, None);
|
|
assert!((scheduler.learning_rate() - 0.01).abs() < 1e-8);
|
|
scheduler.step(9, None);
|
|
assert!((scheduler.learning_rate() - 0.01).abs() < 1e-8);
|
|
|
|
// Epochs 10-19: LR = 0.001
|
|
scheduler.step(10, None);
|
|
assert!((scheduler.learning_rate() - 0.001).abs() < 1e-8);
|
|
|
|
// Epochs 20-29: LR = 0.0001
|
|
scheduler.step(20, None);
|
|
assert!((scheduler.learning_rate() - 0.0001).abs() < 1e-8);
|
|
}
|
|
|
|
#[test]
|
|
fn test_training_session_creation() {
|
|
let config = TrainingConfig::quick();
|
|
let pde_config = PDEConfig::darcy(16);
|
|
let session = TrainingSession::new(config, pde_config);
|
|
|
|
assert!(!session.is_training());
|
|
}
|
|
|
|
// Integration test - runs actual training (slow)
|
|
#[test]
|
|
#[ignore = "slow integration test"]
|
|
fn test_quick_training() {
|
|
let config = TrainingConfig {
|
|
epochs: 2,
|
|
batch_size: 4,
|
|
learning_rate: 0.001,
|
|
n_train_samples: 8,
|
|
n_val_samples: 4,
|
|
};
|
|
let pde_config = PDEConfig::darcy(16)
|
|
.with_modes(4, 4)
|
|
.with_width(8)
|
|
.with_layers(2);
|
|
|
|
// Explicitly use CpuBackend for test
|
|
let trainer = FnoTrainer::<CpuBackend>::new(config, pde_config);
|
|
let result = trainer.train();
|
|
|
|
assert!(
|
|
result.is_ok(),
|
|
"Training should succeed: {:?}",
|
|
result.err()
|
|
);
|
|
|
|
let final_progress = trainer.get_progress();
|
|
assert_eq!(final_progress.status, TrainingStatus::Complete);
|
|
assert_eq!(final_progress.epoch, 2);
|
|
}
|
|
|
|
#[test]
|
|
fn test_trainer_save_weights_creates_file() {
|
|
use tempfile::tempdir;
|
|
|
|
let _temp_dir = tempdir().expect("Failed to create temp dir");
|
|
let device = CpuDevice::default();
|
|
|
|
let model = FNO2d::<CpuBackend>::new_with_layers(1, 1, 8, 4, 2, &device)
|
|
.expect("Failed to create model");
|
|
|
|
let config = TrainingConfig::quick();
|
|
let pde_config = PDEConfig::darcy(16);
|
|
let trainer = FnoTrainer::<CpuBackend>::with_device(config, pde_config, device);
|
|
|
|
let weights_path = trainer
|
|
.save_weights(&model)
|
|
.expect("Failed to save weights");
|
|
|
|
assert!(
|
|
std::path::Path::new(&weights_path).exists(),
|
|
"Weights file should exist"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_trainer_save_weights_roundtrip() {
|
|
use rtx_neural_operator::weights::load_fno2d_weights;
|
|
use tempfile::tempdir;
|
|
|
|
let _temp_dir = tempdir().expect("Failed to create temp dir");
|
|
let device = CpuDevice::default();
|
|
|
|
let model = FNO2d::<CpuBackend>::new_with_layers(1, 1, 8, 4, 2, &device)
|
|
.expect("Failed to create model");
|
|
|
|
let original_weights = model.to_weights("test_pde");
|
|
|
|
let config = TrainingConfig::quick();
|
|
let pde_config = PDEConfig::darcy(16);
|
|
let trainer = FnoTrainer::<CpuBackend>::with_device(config, pde_config, device.clone());
|
|
|
|
let weights_path = trainer
|
|
.save_weights(&model)
|
|
.expect("Failed to save weights");
|
|
|
|
let loaded_weights = load_fno2d_weights(&weights_path).expect("Failed to load weights");
|
|
|
|
assert_eq!(loaded_weights.config.width, original_weights.config.width);
|
|
assert_eq!(
|
|
loaded_weights.config.n_layers,
|
|
original_weights.config.n_layers
|
|
);
|
|
assert_eq!(
|
|
loaded_weights.config.n_modes,
|
|
original_weights.config.n_modes
|
|
);
|
|
|
|
assert_eq!(
|
|
loaded_weights.lifting_fc1_weight.len(),
|
|
original_weights.lifting_fc1_weight.len()
|
|
);
|
|
for (a, b) in loaded_weights
|
|
.lifting_fc1_weight
|
|
.iter()
|
|
.zip(original_weights.lifting_fc1_weight.iter())
|
|
{
|
|
assert!(
|
|
(a - b).abs() < 1e-6,
|
|
"Lifting fc1 weights mismatch: {} vs {}",
|
|
a,
|
|
b
|
|
);
|
|
}
|
|
|
|
assert_eq!(
|
|
loaded_weights.spectral_weights.len(),
|
|
original_weights.spectral_weights.len()
|
|
);
|
|
for (loaded_sw, orig_sw) in loaded_weights
|
|
.spectral_weights
|
|
.iter()
|
|
.zip(original_weights.spectral_weights.iter())
|
|
{
|
|
assert_eq!(loaded_sw.weights1_real.len(), orig_sw.weights1_real.len());
|
|
for (a, b) in loaded_sw
|
|
.weights1_real
|
|
.iter()
|
|
.zip(orig_sw.weights1_real.iter())
|
|
{
|
|
assert!(
|
|
(a - b).abs() < 1e-6,
|
|
"Spectral weights mismatch: {} vs {}",
|
|
a,
|
|
b
|
|
);
|
|
}
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_trainer_save_weights_contains_expected_tensors() {
|
|
use rtx_hub::safetensors::SafeTensors;
|
|
use tempfile::tempdir;
|
|
|
|
let _temp_dir = tempdir().expect("Failed to create temp dir");
|
|
let device = CpuDevice::default();
|
|
|
|
let model = FNO2d::<CpuBackend>::new_with_layers(1, 1, 8, 4, 2, &device)
|
|
.expect("Failed to create model");
|
|
|
|
let config = TrainingConfig::quick();
|
|
let pde_config = PDEConfig::darcy(16);
|
|
let trainer = FnoTrainer::<CpuBackend>::with_device(config, pde_config, device);
|
|
|
|
let weights_path = trainer
|
|
.save_weights(&model)
|
|
.expect("Failed to save weights");
|
|
|
|
let file_bytes = std::fs::read(&weights_path).expect("Failed to read weights file");
|
|
let safetensors =
|
|
SafeTensors::from_bytes(&file_bytes).expect("Failed to parse SafeTensors");
|
|
|
|
let tensor_names = safetensors.tensor_names();
|
|
|
|
assert!(
|
|
tensor_names.contains(&"lifting.fcs.0.weight"),
|
|
"Should contain lifting.fcs.0.weight"
|
|
);
|
|
assert!(
|
|
tensor_names.contains(&"lifting.fcs.0.bias"),
|
|
"Should contain lifting.fcs.0.bias"
|
|
);
|
|
assert!(
|
|
tensor_names.contains(&"lifting.fcs.1.weight"),
|
|
"Should contain lifting.fcs.1.weight"
|
|
);
|
|
assert!(
|
|
tensor_names.contains(&"lifting.fcs.1.bias"),
|
|
"Should contain lifting.fcs.1.bias"
|
|
);
|
|
|
|
assert!(
|
|
tensor_names.contains(&"spectral_conv.0.weights1_real"),
|
|
"Should contain spectral_conv.0.weights1_real"
|
|
);
|
|
assert!(
|
|
tensor_names.contains(&"spectral_conv.0.weights1_imag"),
|
|
"Should contain spectral_conv.0.weights1_imag"
|
|
);
|
|
assert!(
|
|
tensor_names.contains(&"spectral_conv.1.weights1_real"),
|
|
"Should contain spectral_conv.1.weights1_real"
|
|
);
|
|
|
|
assert!(
|
|
tensor_names.contains(&"conv.0.weight"),
|
|
"Should contain conv.0.weight"
|
|
);
|
|
assert!(
|
|
tensor_names.contains(&"conv.0.bias"),
|
|
"Should contain conv.0.bias"
|
|
);
|
|
|
|
assert!(
|
|
tensor_names.contains(&"projection.0.weight"),
|
|
"Should contain projection.0.weight"
|
|
);
|
|
assert!(
|
|
tensor_names.contains(&"projection.0.bias"),
|
|
"Should contain projection.0.bias"
|
|
);
|
|
assert!(
|
|
tensor_names.contains(&"projection.1.weight"),
|
|
"Should contain projection.1.weight"
|
|
);
|
|
assert!(
|
|
tensor_names.contains(&"projection.1.bias"),
|
|
"Should contain projection.1.bias"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_trainer_save_weights_includes_metadata() {
|
|
use rtx_hub::safetensors::SafeTensors;
|
|
use tempfile::tempdir;
|
|
|
|
let _temp_dir = tempdir().expect("Failed to create temp dir");
|
|
let device = CpuDevice::default();
|
|
|
|
let model = FNO2d::<CpuBackend>::new_with_layers(1, 1, 8, 4, 2, &device)
|
|
.expect("Failed to create model");
|
|
|
|
let config = TrainingConfig::quick();
|
|
let pde_config = PDEConfig::darcy(16)
|
|
.with_modes(4, 4)
|
|
.with_width(8)
|
|
.with_layers(2);
|
|
let trainer = FnoTrainer::<CpuBackend>::with_device(config, pde_config.clone(), device);
|
|
|
|
let weights_path = trainer
|
|
.save_weights(&model)
|
|
.expect("Failed to save weights");
|
|
|
|
let file_bytes = std::fs::read(&weights_path).expect("Failed to read weights file");
|
|
let safetensors =
|
|
SafeTensors::from_bytes(&file_bytes).expect("Failed to parse SafeTensors");
|
|
|
|
let metadata = safetensors.metadata().expect("Should have metadata");
|
|
|
|
assert_eq!(
|
|
metadata.get("model_type"),
|
|
Some(&"FNO2d".to_string()),
|
|
"Should have model_type metadata"
|
|
);
|
|
assert_eq!(
|
|
metadata.get("pde_type"),
|
|
Some(&"darcy_flow".to_string()),
|
|
"Should have pde_type metadata"
|
|
);
|
|
assert_eq!(
|
|
metadata.get("resolution"),
|
|
Some(&"16".to_string()),
|
|
"Should have resolution metadata"
|
|
);
|
|
assert_eq!(
|
|
metadata.get("width"),
|
|
Some(&"8".to_string()),
|
|
"Should have width metadata"
|
|
);
|
|
assert_eq!(
|
|
metadata.get("n_layers"),
|
|
Some(&"2".to_string()),
|
|
"Should have n_layers metadata"
|
|
);
|
|
}
|
|
|
|
// ============================================================================
|
|
// TDD Tests for GPU Acceleration (RED-GREEN-REFACTOR)
|
|
// ============================================================================
|
|
|
|
#[test]
|
|
fn test_device_selection_reports_device_name() {
|
|
// RED: Test that trainer reports device name in initial progress
|
|
let config = TrainingConfig::quick();
|
|
let pde_config = PDEConfig::darcy(16);
|
|
let trainer = FnoTrainer::<CpuBackend>::new(config, pde_config);
|
|
|
|
let initial_progress = trainer.get_progress();
|
|
|
|
// GREEN: Device name should be populated
|
|
assert!(
|
|
!initial_progress.device.is_empty(),
|
|
"Device name should not be empty"
|
|
);
|
|
assert!(
|
|
initial_progress.device.contains("CPU")
|
|
|| initial_progress.device.contains("CUDA")
|
|
|| initial_progress.device.contains("Metal"),
|
|
"Device name should identify backend type: got '{}'",
|
|
initial_progress.device
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_training_progress_preserves_device_info() {
|
|
// RED: Test that device info persists through training updates
|
|
let config = TrainingConfig::quick();
|
|
let pde_config = PDEConfig::darcy(16);
|
|
let trainer = FnoTrainer::<CpuBackend>::new(config, pde_config);
|
|
|
|
let initial_device = trainer.get_progress().device.clone();
|
|
|
|
// Update progress
|
|
trainer.update_progress(|p| {
|
|
p.epoch = 1;
|
|
p.loss = 0.5;
|
|
});
|
|
|
|
let updated_progress = trainer.get_progress();
|
|
|
|
// GREEN: Device name should be unchanged
|
|
assert_eq!(
|
|
updated_progress.device, initial_device,
|
|
"Device name should persist through updates"
|
|
);
|
|
assert_eq!(updated_progress.epoch, 1);
|
|
}
|
|
|
|
#[test]
|
|
fn test_cpu_backend_trainer_creation() {
|
|
// RED: Test explicit CPU backend trainer creation
|
|
let config = TrainingConfig::quick();
|
|
let pde_config = PDEConfig::darcy(16);
|
|
let device = CpuDevice::default();
|
|
|
|
let trainer = FnoTrainer::<CpuBackend>::with_device(config, pde_config, device);
|
|
|
|
// GREEN: Verify trainer was created successfully
|
|
assert!(!trainer.is_training());
|
|
let progress = trainer.get_progress();
|
|
assert!(
|
|
progress.device.contains("CPU"),
|
|
"Device should be CPU, got: {}",
|
|
progress.device
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_training_session_auto_selects_backend() {
|
|
// RED: Test that TrainingSession automatically selects best backend
|
|
let config = TrainingConfig::quick();
|
|
let pde_config = PDEConfig::darcy(16);
|
|
|
|
let session = TrainingSession::new(config, pde_config);
|
|
|
|
// GREEN: Session should be created successfully with some backend
|
|
assert!(!session.is_training());
|
|
let progress = session.get_progress();
|
|
assert!(!progress.device.is_empty(), "Device should be selected");
|
|
}
|
|
|
|
#[test]
|
|
fn test_training_session_with_explicit_cpu_backend() {
|
|
// RED: Test creating session with explicit CPU backend
|
|
let config = TrainingConfig::quick();
|
|
let pde_config = PDEConfig::darcy(16);
|
|
|
|
let session = TrainingSession::with_backend::<CpuBackend>(config, pde_config);
|
|
|
|
// GREEN: Should create session with CPU backend
|
|
assert!(!session.is_training());
|
|
let progress = session.get_progress();
|
|
assert!(progress.device.contains("CPU"), "Should use CPU backend");
|
|
}
|
|
|
|
#[test]
|
|
#[cfg(feature = "cuda")]
|
|
fn test_training_session_with_cuda_backend() {
|
|
// RED: Test creating session with CUDA backend if available
|
|
let config = TrainingConfig::quick();
|
|
let pde_config = PDEConfig::darcy(16);
|
|
|
|
// Try to create CUDA device
|
|
if CudaDevice::default_available() {
|
|
let session = TrainingSession::with_backend::<CudaBackend>(config, pde_config);
|
|
|
|
// GREEN: Should create session with CUDA backend
|
|
assert!(!session.is_training());
|
|
let progress = session.get_progress();
|
|
assert!(
|
|
progress.device.contains("CUDA"),
|
|
"Should use CUDA backend, got: {}",
|
|
progress.device
|
|
);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
#[cfg(feature = "metal")]
|
|
fn test_training_session_with_metal_backend() {
|
|
// RED: Test creating session with Metal backend if available
|
|
let config = TrainingConfig::quick();
|
|
let pde_config = PDEConfig::darcy(16);
|
|
|
|
// Try to create Metal device
|
|
use rtx_backend_metal::MetalDevice;
|
|
if MetalDevice::default_available() {
|
|
let session = TrainingSession::with_backend::<MetalBackend>(config, pde_config);
|
|
|
|
// GREEN: Should create session with Metal backend
|
|
assert!(!session.is_training());
|
|
let progress = session.get_progress();
|
|
assert!(
|
|
progress.device.contains("Metal"),
|
|
"Should use Metal backend, got: {}",
|
|
progress.device
|
|
);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_get_device_string_returns_valid_name() {
|
|
// RED: Test that get_device_string returns a valid device identifier
|
|
let device_string = get_device_string();
|
|
|
|
// GREEN: Should return non-empty string with backend type
|
|
assert!(
|
|
!device_string.is_empty(),
|
|
"Device string should not be empty"
|
|
);
|
|
assert!(
|
|
device_string.contains("CPU")
|
|
|| device_string.contains("CUDA")
|
|
|| device_string.contains("Metal")
|
|
|| device_string.contains("ROCm"),
|
|
"Device string should contain backend type: got '{}'",
|
|
device_string
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_initial_progress_includes_learning_rate() {
|
|
// RED: Test that initial progress includes the configured learning rate
|
|
let config = TrainingConfig {
|
|
epochs: 10,
|
|
batch_size: 8,
|
|
learning_rate: 0.002,
|
|
n_train_samples: 100,
|
|
n_val_samples: 20,
|
|
};
|
|
let pde_config = PDEConfig::darcy(16);
|
|
let trainer = FnoTrainer::<CpuBackend>::new(config.clone(), pde_config);
|
|
|
|
let initial_progress = trainer.get_progress();
|
|
|
|
// GREEN: Initial LR should match config
|
|
assert!(
|
|
(initial_progress.current_lr - config.learning_rate).abs() < 1e-6,
|
|
"Initial LR should be {}, got {}",
|
|
config.learning_rate,
|
|
initial_progress.current_lr
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_backend_selection_order() {
|
|
// RED: Test that select_best_backend returns available backend
|
|
use rtx_backend::auto_select::select_best_backend;
|
|
|
|
let backend = select_best_backend();
|
|
|
|
// GREEN: Should return one of the valid backend types (case insensitive)
|
|
let backend_str = backend.to_string();
|
|
let backend_lower = backend_str.to_lowercase();
|
|
assert!(
|
|
backend_lower == "cpu"
|
|
|| backend_lower == "cuda"
|
|
|| backend_lower == "metal"
|
|
|| backend_lower == "rocm",
|
|
"Backend selection returned unexpected type: {}",
|
|
backend_str
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_dynamic_trainer_trait_implementation() {
|
|
// RED: Test that FnoTrainer implements DynamicTrainer
|
|
let config = TrainingConfig::quick();
|
|
let pde_config = PDEConfig::darcy(16);
|
|
let trainer = FnoTrainer::<CpuBackend>::new(config, pde_config);
|
|
|
|
let dynamic: &dyn DynamicTrainer = &trainer;
|
|
|
|
// GREEN: Should be able to call trait methods
|
|
assert!(!dynamic.is_training());
|
|
let progress = dynamic.get_progress();
|
|
assert_eq!(progress.status, TrainingStatus::NotStarted);
|
|
}
|
|
|
|
#[test]
|
|
fn test_training_session_creation_with_auto_backend() {
|
|
// RED: Test that new() uses automatic backend selection
|
|
let config = TrainingConfig::quick();
|
|
let pde_config = PDEConfig::darcy(16);
|
|
|
|
let session = TrainingSession::new(config, pde_config);
|
|
let progress = session.get_progress();
|
|
|
|
// GREEN: Should have selected some backend
|
|
assert!(!progress.device.is_empty());
|
|
|
|
// Device string should match one of the known formats
|
|
let device_lower = progress.device.to_lowercase();
|
|
assert!(
|
|
device_lower.contains("cpu")
|
|
|| device_lower.contains("cuda")
|
|
|| device_lower.contains("metal")
|
|
|| device_lower.contains("rocm"),
|
|
"Device should be recognized backend type: {}",
|
|
progress.device
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_progress_device_field_serialization() {
|
|
// RED: Test that TrainingProgress with device field serializes correctly
|
|
// Create a complete progress object with non-infinite values
|
|
let mut progress = TrainingProgress::new(10, 5);
|
|
progress.device = "CUDA (RTX 4090)".to_string();
|
|
progress.current_lr = 0.001;
|
|
progress.loss = 0.5;
|
|
progress.best_loss = 0.3;
|
|
progress.epoch = 5;
|
|
progress.batch = 2;
|
|
|
|
let json = serde_json::to_string(&progress).expect("Should serialize");
|
|
let deserialized: TrainingProgress =
|
|
serde_json::from_str(&json).expect("Should deserialize");
|
|
|
|
// GREEN: Device field and learning rate should roundtrip correctly
|
|
assert_eq!(deserialized.device, "CUDA (RTX 4090)");
|
|
assert!((deserialized.current_lr - 0.001).abs() < 1e-6);
|
|
assert!((deserialized.loss - 0.5).abs() < 1e-6);
|
|
assert!((deserialized.best_loss - 0.3).abs() < 1e-6);
|
|
assert_eq!(deserialized.epoch, 5);
|
|
assert_eq!(deserialized.batch, 2);
|
|
}
|
|
|
|
#[test]
|
|
fn test_trainer_reports_correct_device_for_cpu() {
|
|
// RED: Test that CPU backend reports CPU device
|
|
let config = TrainingConfig::quick();
|
|
let pde_config = PDEConfig::darcy(16);
|
|
let device = CpuDevice::default();
|
|
let trainer = FnoTrainer::<CpuBackend>::with_device(config, pde_config, device);
|
|
|
|
let progress = trainer.get_progress();
|
|
|
|
// GREEN: Device should indicate CPU
|
|
assert!(
|
|
progress.device.contains("CPU"),
|
|
"CPU backend should report CPU device, got: {}",
|
|
progress.device
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_multiple_trainers_with_different_backends() {
|
|
// RED: Test creating multiple trainers with different backends
|
|
let config1 = TrainingConfig::quick();
|
|
let config2 = TrainingConfig::quick();
|
|
let pde_config1 = PDEConfig::darcy(16);
|
|
let pde_config2 = PDEConfig::darcy(16);
|
|
|
|
let trainer1 = FnoTrainer::<CpuBackend>::new(config1, pde_config1);
|
|
let session2 = TrainingSession::new(config2, pde_config2);
|
|
|
|
// GREEN: Both should be created successfully
|
|
assert!(!trainer1.is_training());
|
|
assert!(!session2.is_training());
|
|
|
|
let progress1 = trainer1.get_progress();
|
|
let progress2 = session2.get_progress();
|
|
|
|
assert!(!progress1.device.is_empty());
|
|
assert!(!progress2.device.is_empty());
|
|
}
|
|
}
|