Files
rustytorch/demos/rtx-pinn-benchmark/src/training.rs
T
2026-03-04 00:08:42 +00:00

596 lines
19 KiB
Rust

//! Training module with AdamW optimizer
use crate::error::{PINNError, Result};
/// AdamW optimizer state
#[derive(Debug, Clone)]
pub struct AdamW {
/// Learning rate
learning_rate: f64,
/// Beta1 (exponential decay for first moment)
beta1: f64,
/// Beta2 (exponential decay for second moment)
beta2: f64,
/// Weight decay coefficient
weight_decay: f64,
/// Epsilon for numerical stability
epsilon: f64,
/// Time step counter
t: usize,
/// First moment estimates for weights
m_weights: Vec<Vec<Vec<f64>>>,
/// First moment estimates for biases
m_biases: Vec<Vec<f64>>,
/// Second moment estimates for weights
v_weights: Vec<Vec<Vec<f64>>>,
/// Second moment estimates for biases
v_biases: Vec<Vec<f64>>,
}
impl AdamW {
/// Creates a new AdamW optimizer
///
/// # Arguments
///
/// * `learning_rate` - Learning rate (typically 0.001)
/// * `beta1` - Exponential decay rate for first moment (typically 0.9)
/// * `beta2` - Exponential decay rate for second moment (typically 0.999)
/// * `weight_decay` - Weight decay coefficient (typically 0.01)
/// * `weight_shapes` - Shapes of weight matrices
/// * `bias_shapes` - Shapes of bias vectors
///
/// # Errors
///
/// Returns an error if hyperparameters are invalid
pub fn new(
learning_rate: f64,
beta1: f64,
beta2: f64,
weight_decay: f64,
weight_shapes: &[Vec<Vec<f64>>],
bias_shapes: &[Vec<f64>],
) -> Result<Self> {
if learning_rate <= 0.0 {
return Err(PINNError::training("learning rate must be positive"));
}
if !(0.0..1.0).contains(&beta1) {
return Err(PINNError::training("beta1 must be in [0, 1)"));
}
if !(0.0..1.0).contains(&beta2) {
return Err(PINNError::training("beta2 must be in [0, 1)"));
}
if weight_decay < 0.0 {
return Err(PINNError::training("weight decay must be non-negative"));
}
// Initialize moment estimates with zeros
let m_weights: Vec<Vec<Vec<f64>>> = weight_shapes
.iter()
.map(|layer| layer.iter().map(|row| vec![0.0; row.len()]).collect())
.collect();
let m_biases: Vec<Vec<f64>> = bias_shapes
.iter()
.map(|layer| vec![0.0; layer.len()])
.collect();
let v_weights = m_weights.clone();
let v_biases = m_biases.clone();
Ok(Self {
learning_rate,
beta1,
beta2,
weight_decay,
epsilon: 1e-8,
t: 0,
m_weights,
m_biases,
v_weights,
v_biases,
})
}
/// Performs one optimization step
///
/// # Arguments
///
/// * `weight_grads` - Gradients for weights
/// * `bias_grads` - Gradients for biases
///
/// # Returns
///
/// Weight and bias updates to apply to the network
pub fn step(
&mut self,
weight_grads: &[Vec<Vec<f64>>],
bias_grads: &[Vec<f64>],
) -> (Vec<Vec<Vec<f64>>>, Vec<Vec<f64>>) {
self.t += 1;
// Bias correction terms
let bias_correction1 = 1.0 - self.beta1.powi(self.t as i32);
let bias_correction2 = 1.0 - self.beta2.powi(self.t as i32);
let mut weight_updates = Vec::new();
let mut bias_updates = Vec::new();
// Update weights
for (layer_idx, grads) in weight_grads.iter().enumerate() {
let mut layer_updates = Vec::new();
for (i, grad_row) in grads.iter().enumerate() {
let mut row_updates = Vec::new();
for (j, &grad) in grad_row.iter().enumerate() {
// Update biased first moment estimate
self.m_weights[layer_idx][i][j] =
self.beta1 * self.m_weights[layer_idx][i][j] + (1.0 - self.beta1) * grad;
// Update biased second moment estimate
self.v_weights[layer_idx][i][j] = self.beta2 * self.v_weights[layer_idx][i][j]
+ (1.0 - self.beta2) * grad * grad;
// Compute bias-corrected moment estimates
let m_hat = self.m_weights[layer_idx][i][j] / bias_correction1;
let v_hat = self.v_weights[layer_idx][i][j] / bias_correction2;
// Compute update (negative because we want to minimize)
let update = -self.learning_rate * (m_hat / (v_hat.sqrt() + self.epsilon));
row_updates.push(update);
}
layer_updates.push(row_updates);
}
weight_updates.push(layer_updates);
}
// Update biases
for (layer_idx, grads) in bias_grads.iter().enumerate() {
let mut layer_updates = Vec::new();
for (i, &grad) in grads.iter().enumerate() {
// Update biased first moment estimate
self.m_biases[layer_idx][i] =
self.beta1 * self.m_biases[layer_idx][i] + (1.0 - self.beta1) * grad;
// Update biased second moment estimate
self.v_biases[layer_idx][i] =
self.beta2 * self.v_biases[layer_idx][i] + (1.0 - self.beta2) * grad * grad;
// Compute bias-corrected moment estimates
let m_hat = self.m_biases[layer_idx][i] / bias_correction1;
let v_hat = self.v_biases[layer_idx][i] / bias_correction2;
// Compute update
let update = -self.learning_rate * (m_hat / (v_hat.sqrt() + self.epsilon));
layer_updates.push(update);
}
bias_updates.push(layer_updates);
}
(weight_updates, bias_updates)
}
/// Gets the current learning rate
pub const fn learning_rate(&self) -> f64 {
self.learning_rate
}
/// Sets a new learning rate
pub fn set_learning_rate(&mut self, lr: f64) {
self.learning_rate = lr;
}
/// Gets the current time step
pub const fn time_step(&self) -> usize {
self.t
}
}
/// Training configuration
#[derive(Debug, Clone)]
pub struct TrainingConfig {
/// Number of training epochs
pub num_epochs: usize,
/// Number of collocation points
pub num_collocation_points: usize,
/// Number of boundary points
pub num_boundary_points: usize,
/// Learning rate
pub learning_rate: f64,
/// Weight for physics loss
pub physics_weight: f64,
/// Weight for boundary loss
pub boundary_weight: f64,
}
impl TrainingConfig {
/// Creates a new training configuration with validation
///
/// # Errors
///
/// Returns an error if configuration is invalid
pub fn new(
num_epochs: usize,
num_collocation_points: usize,
num_boundary_points: usize,
learning_rate: f64,
physics_weight: f64,
boundary_weight: f64,
) -> Result<Self> {
if num_epochs == 0 {
return Err(PINNError::invalid_config("num_epochs must be positive"));
}
if num_collocation_points == 0 {
return Err(PINNError::invalid_config(
"num_collocation_points must be positive",
));
}
if num_boundary_points == 0 {
return Err(PINNError::invalid_config(
"num_boundary_points must be positive",
));
}
if learning_rate <= 0.0 {
return Err(PINNError::invalid_config("learning_rate must be positive"));
}
if physics_weight < 0.0 || boundary_weight < 0.0 {
return Err(PINNError::invalid_config(
"loss weights must be non-negative",
));
}
Ok(Self {
num_epochs,
num_collocation_points,
num_boundary_points,
learning_rate,
physics_weight,
boundary_weight,
})
}
}
/// Computes mean squared error loss
pub fn mse_loss(predictions: &[f64], targets: &[f64]) -> Result<f64> {
if predictions.len() != targets.len() {
return Err(PINNError::training(
"predictions and targets length mismatch",
));
}
if predictions.is_empty() {
return Err(PINNError::training("empty predictions"));
}
let sum_squared_error: f64 = predictions
.iter()
.zip(targets.iter())
.map(|(pred, target)| (pred - target).powi(2))
.sum();
Ok(sum_squared_error / predictions.len() as f64)
}
#[cfg(test)]
mod tests {
use super::*;
use approx::assert_abs_diff_eq;
#[test]
fn test_adamw_creation_valid() {
let weight_shapes = vec![vec![vec![0.0; 2]; 3], vec![vec![0.0; 3]; 1]];
let bias_shapes = vec![vec![0.0; 3], vec![0.0; 1]];
let optimizer = AdamW::new(0.001, 0.9, 0.999, 0.01, &weight_shapes, &bias_shapes);
assert!(optimizer.is_ok());
let optimizer = optimizer.unwrap();
assert_abs_diff_eq!(optimizer.learning_rate(), 0.001);
assert_eq!(optimizer.time_step(), 0);
}
#[test]
fn test_adamw_creation_invalid_lr() {
let weight_shapes = vec![vec![vec![0.0; 2]; 3]];
let bias_shapes = vec![vec![0.0; 3]];
let optimizer = AdamW::new(0.0, 0.9, 0.999, 0.01, &weight_shapes, &bias_shapes);
assert!(optimizer.is_err());
let optimizer = AdamW::new(-0.001, 0.9, 0.999, 0.01, &weight_shapes, &bias_shapes);
assert!(optimizer.is_err());
}
#[test]
fn test_adamw_creation_invalid_beta1() {
let weight_shapes = vec![vec![vec![0.0; 2]; 3]];
let bias_shapes = vec![vec![0.0; 3]];
let optimizer = AdamW::new(0.001, 1.0, 0.999, 0.01, &weight_shapes, &bias_shapes);
assert!(optimizer.is_err());
let optimizer = AdamW::new(0.001, -0.1, 0.999, 0.01, &weight_shapes, &bias_shapes);
assert!(optimizer.is_err());
}
#[test]
fn test_adamw_creation_invalid_beta2() {
let weight_shapes = vec![vec![vec![0.0; 2]; 3]];
let bias_shapes = vec![vec![0.0; 3]];
let optimizer = AdamW::new(0.001, 0.9, 1.0, 0.01, &weight_shapes, &bias_shapes);
assert!(optimizer.is_err());
}
#[test]
fn test_adamw_creation_invalid_weight_decay() {
let weight_shapes = vec![vec![vec![0.0; 2]; 3]];
let bias_shapes = vec![vec![0.0; 3]];
let optimizer = AdamW::new(0.001, 0.9, 0.999, -0.01, &weight_shapes, &bias_shapes);
assert!(optimizer.is_err());
}
#[test]
fn test_adamw_step_increments_time() {
let weight_shapes = vec![vec![vec![0.0; 2]; 3]];
let bias_shapes = vec![vec![0.0; 3]];
let mut optimizer =
AdamW::new(0.001, 0.9, 0.999, 0.01, &weight_shapes, &bias_shapes).unwrap();
assert_eq!(optimizer.time_step(), 0);
let weight_grads = vec![vec![vec![0.1; 2]; 3]];
let bias_grads = vec![vec![0.1; 3]];
optimizer.step(&weight_grads, &bias_grads);
assert_eq!(optimizer.time_step(), 1);
optimizer.step(&weight_grads, &bias_grads);
assert_eq!(optimizer.time_step(), 2);
}
#[test]
fn test_adamw_step_returns_updates() {
let weight_shapes = vec![vec![vec![0.0; 2]; 1]];
let bias_shapes = vec![vec![0.0; 1]];
let mut optimizer =
AdamW::new(0.001, 0.9, 0.999, 0.0, &weight_shapes, &bias_shapes).unwrap();
let weight_grads = vec![vec![vec![1.0; 2]; 1]];
let bias_grads = vec![vec![1.0; 1]];
let (weight_updates, bias_updates) = optimizer.step(&weight_grads, &bias_grads);
assert_eq!(weight_updates.len(), 1);
assert_eq!(weight_updates[0].len(), 1);
assert_eq!(weight_updates[0][0].len(), 2);
assert_eq!(bias_updates.len(), 1);
assert_eq!(bias_updates[0].len(), 1);
}
#[test]
fn test_adamw_step_update_sign() {
let weight_shapes = vec![vec![vec![0.0; 1]; 1]];
let bias_shapes = vec![vec![0.0; 1]];
let mut optimizer =
AdamW::new(0.001, 0.9, 0.999, 0.0, &weight_shapes, &bias_shapes).unwrap();
// Positive gradient should give negative update (gradient descent)
let weight_grads = vec![vec![vec![1.0]; 1]];
let bias_grads = vec![vec![1.0]];
let (weight_updates, _) = optimizer.step(&weight_grads, &bias_grads);
assert!(weight_updates[0][0][0] < 0.0);
}
#[test]
fn test_adamw_bias_correction() {
let weight_shapes = vec![vec![vec![0.0; 1]; 1]];
let bias_shapes = vec![vec![0.0; 1]];
let mut optimizer =
AdamW::new(0.001, 0.9, 0.999, 0.0, &weight_shapes, &bias_shapes).unwrap();
let weight_grads = vec![vec![vec![1.0]; 1]];
let bias_grads = vec![vec![1.0]];
let (updates1, _) = optimizer.step(&weight_grads, &bias_grads);
let (updates2, _) = optimizer.step(&weight_grads, &bias_grads);
// First step should have larger magnitude due to bias correction
assert!(updates1[0][0][0].abs() > updates2[0][0][0].abs());
}
#[test]
fn test_adamw_learning_rate_setter() {
let weight_shapes = vec![vec![vec![0.0; 1]; 1]];
let bias_shapes = vec![vec![0.0; 1]];
let mut optimizer =
AdamW::new(0.001, 0.9, 0.999, 0.0, &weight_shapes, &bias_shapes).unwrap();
assert_abs_diff_eq!(optimizer.learning_rate(), 0.001);
optimizer.set_learning_rate(0.0005);
assert_abs_diff_eq!(optimizer.learning_rate(), 0.0005);
}
#[test]
fn test_training_config_creation_valid() {
let config = TrainingConfig::new(1000, 10000, 100, 0.001, 1.0, 1.0);
assert!(config.is_ok());
let config = config.unwrap();
assert_eq!(config.num_epochs, 1000);
assert_eq!(config.num_collocation_points, 10000);
assert_eq!(config.num_boundary_points, 100);
assert_abs_diff_eq!(config.learning_rate, 0.001);
}
#[test]
fn test_training_config_zero_epochs() {
let config = TrainingConfig::new(0, 10000, 100, 0.001, 1.0, 1.0);
assert!(config.is_err());
}
#[test]
fn test_training_config_zero_collocation_points() {
let config = TrainingConfig::new(1000, 0, 100, 0.001, 1.0, 1.0);
assert!(config.is_err());
}
#[test]
fn test_training_config_zero_boundary_points() {
let config = TrainingConfig::new(1000, 10000, 0, 0.001, 1.0, 1.0);
assert!(config.is_err());
}
#[test]
fn test_training_config_invalid_learning_rate() {
let config = TrainingConfig::new(1000, 10000, 100, 0.0, 1.0, 1.0);
assert!(config.is_err());
let config = TrainingConfig::new(1000, 10000, 100, -0.001, 1.0, 1.0);
assert!(config.is_err());
}
#[test]
fn test_training_config_negative_weights() {
let config = TrainingConfig::new(1000, 10000, 100, 0.001, -1.0, 1.0);
assert!(config.is_err());
let config = TrainingConfig::new(1000, 10000, 100, 0.001, 1.0, -1.0);
assert!(config.is_err());
}
#[test]
fn test_mse_loss_perfect_predictions() {
let predictions = vec![1.0, 2.0, 3.0, 4.0];
let targets = vec![1.0, 2.0, 3.0, 4.0];
let loss = mse_loss(&predictions, &targets).unwrap();
assert_abs_diff_eq!(loss, 0.0, epsilon = 1e-10);
}
#[test]
fn test_mse_loss_calculation() {
let predictions = vec![1.0, 2.0, 3.0];
let targets = vec![1.1, 2.1, 2.9];
let loss = mse_loss(&predictions, &targets).unwrap();
// MSE = ((0.1)^2 + (0.1)^2 + (0.1)^2) / 3 = 0.03 / 3 = 0.01
assert_abs_diff_eq!(loss, 0.01, epsilon = 1e-10);
}
#[test]
fn test_mse_loss_length_mismatch() {
let predictions = vec![1.0, 2.0];
let targets = vec![1.0, 2.0, 3.0];
let result = mse_loss(&predictions, &targets);
assert!(result.is_err());
}
#[test]
fn test_mse_loss_empty_arrays() {
let predictions: Vec<f64> = vec![];
let targets: Vec<f64> = vec![];
let result = mse_loss(&predictions, &targets);
assert!(result.is_err());
}
#[test]
fn test_mse_loss_single_value() {
let predictions = vec![5.0];
let targets = vec![3.0];
let loss = mse_loss(&predictions, &targets).unwrap();
// MSE = (5-3)^2 = 4
assert_abs_diff_eq!(loss, 4.0, epsilon = 1e-10);
}
#[test]
fn test_mse_loss_large_errors() {
let predictions = vec![10.0, 20.0];
let targets = vec![0.0, 0.0];
let loss = mse_loss(&predictions, &targets).unwrap();
// MSE = (100 + 400) / 2 = 250
assert_abs_diff_eq!(loss, 250.0, epsilon = 1e-10);
}
#[test]
fn test_adamw_multiple_layers() {
let weight_shapes = vec![
vec![vec![0.0; 2]; 5],
vec![vec![0.0; 5]; 3],
vec![vec![0.0; 3]; 1],
];
let bias_shapes = vec![vec![0.0; 5], vec![0.0; 3], vec![0.0; 1]];
let mut optimizer =
AdamW::new(0.001, 0.9, 0.999, 0.01, &weight_shapes, &bias_shapes).unwrap();
let weight_grads = vec![
vec![vec![0.5; 2]; 5],
vec![vec![0.3; 5]; 3],
vec![vec![0.1; 3]; 1],
];
let bias_grads = vec![vec![0.5; 5], vec![0.3; 3], vec![0.1; 1]];
let (weight_updates, bias_updates) = optimizer.step(&weight_grads, &bias_grads);
assert_eq!(weight_updates.len(), 3);
assert_eq!(bias_updates.len(), 3);
// Check shapes match
for (i, updates) in weight_updates.iter().enumerate() {
assert_eq!(updates.len(), weight_shapes[i].len());
for (j, row) in updates.iter().enumerate() {
assert_eq!(row.len(), weight_shapes[i][j].len());
}
}
}
#[test]
fn test_adamw_momentum_accumulation() {
let weight_shapes = vec![vec![vec![0.0; 1]; 1]];
let bias_shapes = vec![vec![0.0; 1]];
let mut optimizer =
AdamW::new(0.001, 0.9, 0.999, 0.0, &weight_shapes, &bias_shapes).unwrap();
let weight_grads = vec![vec![vec![1.0]; 1]];
let bias_grads = vec![vec![1.0]];
// Take multiple steps with same gradient
let (updates1, _) = optimizer.step(&weight_grads, &bias_grads);
let (updates2, _) = optimizer.step(&weight_grads, &bias_grads);
let (updates3, _) = optimizer.step(&weight_grads, &bias_grads);
// Update magnitudes should stabilize (not grow indefinitely)
let mag1 = updates1[0][0][0].abs();
let mag2 = updates2[0][0][0].abs();
let mag3 = updates3[0][0][0].abs();
assert!(mag2 < mag1); // Second step smaller due to bias correction
assert!(mag3 < mag2 || (mag3 - mag2).abs() < 0.0001); // Should stabilize
}
}