use super::{Optimizer, OptimizerConfig, OptimizerConfigClone};
use crate::{common::matrix::DMat, error::NetworkError, LearningRateScheduler};
use serde::{Deserialize, Serialize};
use typetag;
#[derive(Serialize, Deserialize, Clone)]
struct MomentumOptimizer {
config: MomentumConfig,
velocity_weights: DMat,
velocity_biases: DMat,
}
impl MomentumOptimizer {
pub fn new(config: MomentumConfig) -> Self {
Self {
config,
velocity_weights: DMat::zeros(0, 0),
velocity_biases: DMat::zeros(0, 0),
}
}
}
#[derive(Serialize, Deserialize, Clone)]
struct MomentumConfig {
learning_rate: f32,
momentum: f32,
scheduler: Option<Box<dyn LearningRateScheduler>>,
}
#[typetag::serde]
impl OptimizerConfig for MomentumConfig {
fn update_learning_rate(&mut self, learning_rate: f32) {
self.learning_rate = learning_rate;
}
fn create_optimizer(&self) -> Box<dyn Optimizer> {
Box::new(MomentumOptimizer::new(self.clone()))
}
fn learning_rate(&self) -> f32 {
self.learning_rate
}
}
#[typetag::serde]
impl Optimizer for MomentumOptimizer {
fn initialize(&mut self, weights: &DMat, biases: &DMat) {
self.velocity_weights = DMat::zeros(weights.rows(), weights.cols());
self.velocity_biases = DMat::zeros(biases.rows(), biases.cols());
}
fn update(&mut self, weights: &mut DMat, biases: &mut DMat, d_weights: &DMat, d_biases: &DMat, epoch: usize) {
if self.config.scheduler.is_some() {
let scheduler = self.config.scheduler.as_ref().unwrap();
self.config.learning_rate = scheduler.schedule(epoch, self.config.learning_rate);
}
weights.apply_with_indices(|r, c, v| {
let velocity = &mut self.velocity_weights;
let previous_velocity = velocity.at(r, c);
let new_velocity =
self.config.momentum * previous_velocity + self.config.learning_rate * d_weights.at(r, c);
velocity.set(r, c, new_velocity);
*v -= new_velocity
});
biases.apply_with_indices(|r, c, v| {
let velocity = &mut self.velocity_biases;
let previous_velocity = velocity.at(r, c);
let new_velocity = self.config.momentum * previous_velocity + self.config.learning_rate * d_biases.at(r, c);
velocity.set(r, c, new_velocity);
*v -= new_velocity
});
}
fn update_learning_rate(&mut self, learning_rate: f32) {
self.config.learning_rate = learning_rate;
}
}
pub struct Momentum {
learning_rate: f32,
momentum: f32,
scheduler: Option<Result<Box<dyn LearningRateScheduler>, NetworkError>>,
}
impl Momentum {
fn new() -> Self {
Self {
learning_rate: 0.01,
momentum: 0.9,
scheduler: None,
}
}
pub fn learning_rate(mut self, lr: f32) -> Self {
self.learning_rate = lr;
self
}
pub fn momentum(mut self, momentum: f32) -> Self {
self.momentum = momentum;
self
}
pub fn scheduler(mut self, scheduler: Result<Box<dyn LearningRateScheduler>, NetworkError>) -> Self {
self.scheduler = Some(scheduler);
self
}
fn validate(&self) -> Result<(), NetworkError> {
if self.learning_rate <= 0.0 {
return Err(NetworkError::ConfigError(format!(
"Learning rate for Momentum must be greater than 0.0, but was {}",
self.learning_rate
)));
}
if self.momentum < 0.0 || self.momentum > 1.0 {
return Err(NetworkError::ConfigError(format!(
"Momentum for Momentum must be in [0.0, 1.0], but was {}",
self.momentum
)));
}
if let Some(ref scheduler) = self.scheduler {
scheduler.as_ref().map_err(|e| e.clone())?;
}
Ok(())
}
pub fn build(self) -> Result<Box<dyn OptimizerConfig>, NetworkError> {
self.validate()?;
Ok(Box::new(MomentumConfig {
learning_rate: self.learning_rate,
momentum: self.momentum,
scheduler: self.scheduler.map(|s| s.unwrap()),
}))
}
}
impl Default for Momentum {
fn default() -> Self {
Self::new()
}
}
impl OptimizerConfigClone for MomentumConfig {
fn clone_box(&self) -> Box<dyn OptimizerConfig> {
Box::new(self.clone())
}
}
#[cfg(test)]
mod tests {
use crate::{step::Step, util::equal_approx};
use super::*;
#[test]
fn test_initialize() {
let config = MomentumConfig {
learning_rate: 0.01,
momentum: 0.9,
scheduler: None,
};
let mut optimizer = MomentumOptimizer::new(config);
let weights = DMat::new(2, 2, &[0.1, 0.2, 0.3, 0.4]);
let biases = DMat::new(2, 1, &[0.1, 0.2]);
optimizer.initialize(&weights, &biases);
assert_eq!(optimizer.velocity_weights.rows(), 2);
assert_eq!(optimizer.velocity_weights.cols(), 2);
assert_eq!(optimizer.velocity_biases.rows(), 2);
assert_eq!(optimizer.velocity_biases.cols(), 1);
}
#[test]
fn test_update() {
let config = MomentumConfig {
learning_rate: 0.1,
momentum: 0.9,
scheduler: None,
};
let mut optimizer = MomentumOptimizer::new(config);
let mut weights = DMat::new(2, 2, &[1.0, 1.0, 1.0, 1.0]);
let mut biases = DMat::new(2, 1, &[1.0, 1.0]);
let d_weights = DMat::new(2, 2, &[0.1, 0.1, 0.1, 0.1]);
let d_biases = DMat::new(2, 1, &[0.1, 0.1]);
optimizer.initialize(&weights, &biases);
optimizer.update(&mut weights, &mut biases, &d_weights, &d_biases, 1);
assert!(weights.at(0, 0) < 1.0);
assert!(biases.at(0, 0) < 1.0);
}
#[test]
fn test_update_learning_rate() {
let config = MomentumConfig {
learning_rate: 0.01,
momentum: 0.9,
scheduler: None,
};
let mut optimizer = MomentumOptimizer::new(config);
optimizer.update_learning_rate(0.02);
assert_eq!(optimizer.config.learning_rate, 0.02);
}
#[test]
fn test_momentum_optimizer() {
let mut weights = DMat::new(2, 2, &[1.0, 2.0, 3.0, 4.0]);
let mut biases = DMat::new(2, 1, &[1.0, 2.0]);
let d_weights = DMat::new(2, 2, &[0.1, 0.1, 0.1, 0.1]);
let d_biases = DMat::new(2, 1, &[0.1, 0.1]);
let config = MomentumConfig {
learning_rate: 0.1,
momentum: 0.9,
scheduler: None,
};
let mut optimizer = MomentumOptimizer::new(config);
optimizer.initialize(&weights, &biases);
optimizer.update(&mut weights, &mut biases, &d_weights, &d_biases, 1);
let expected_weights = DMat::new(2, 2, &[0.99, 1.99, 2.99, 3.99]);
assert!(equal_approx(&weights, &expected_weights, 1e-3));
}
#[test]
fn test_momentum_builder() {
let optimizer = Momentum::default().learning_rate(0.01).momentum(0.9).build().unwrap();
assert_eq!(optimizer.learning_rate(), 0.01);
}
#[test]
fn test_momentum_builder_invalid_learning_rate() {
let result = Momentum::default().learning_rate(-0.01).build();
assert!(result.is_err());
if let Err(e) = result {
assert_eq!(
e.to_string(),
"Configuration error: Learning rate for Momentum must be greater than 0.0, but was -0.01"
);
}
}
#[test]
fn test_momentum_builder_invalid_scheduler() {
let result = Momentum::new()
.momentum(0.5)
.scheduler(Step::default().decay_rate(5.0).build())
.build();
assert!(result.is_err());
if let Err(e) = result {
assert_eq!(
e.to_string(),
"Configuration error: Decay rate for Step must be in the range (0, 1), but was 5"
);
}
}
}