rustyml 0.17.0

Machine learning models, neural networks, metrics, and data utilities in Rust
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122

//! AdamW optimizer: Adam with decoupled weight decay
//!
//! Holds the [`AdamW`] struct, its constructor, its [`AdamW::with_global_clipnorm`] builder method,
//! and its [`Optimizer`] implementation. The moment math and the update loop live in the private
//! `AdamCore` type, shared with [`Adam`]
//!
//! [`Optimizer`]: crate::neural_network::traits::Optimizer
//! [`Adam`]: crate::neural_network::optimizers::Adam

use crate::error::Error;
use crate::neural_network::LayerPath;
use crate::neural_network::ctx::Grads;
use crate::neural_network::optimizers::adam_core::AdamCore;
use crate::neural_network::traits::{LayerBase, Optimizer};

/// AdamW (Adam with decoupled weight decay) optimizer
///
/// Identical adaptive moment math to [`Adam`](crate::neural_network::optimizers::Adam), but its
/// `weight_decay` is decoupled. AdamW multiplies the parameter by `(1 - learning_rate * weight_decay)` before the gradient
/// step. It adds no L2 term to the gradient. The decay therefore does not flow through the moment estimates, and
/// the adaptive denominator does not rescale it. This follows the Loshchilov and Hutter formulation. With `weight_decay == 0.0` it matches
/// `Adam`.
///
/// Weight decay applies to weight tensors only, never to biases or normalization
/// scale/shift parameters
#[derive(Debug)]
pub struct AdamW {
    /// The shared Adam-family optimizer state
    core: AdamCore,
}

impl AdamW {
    /// Creates a new AdamW optimizer with the specified hyperparameters
    ///
    /// Validates hyperparameters and initializes internal timestep tracking
    ///
    /// # Parameters
    ///
    /// - `learning_rate` - step size for parameter updates
    /// - `beta1` - decay rate for the first moment estimates (typically 0.9)
    /// - `beta2` - decay rate for the second moment estimates (typically 0.999)
    /// - `epsilon` - small constant for numerical stability (typically 1e-8)
    /// - `weight_decay` - decoupled weight-decay coefficient applied directly to the parameters.
    ///   `0.0` disables it. For classic coupled L2 decay use
    ///   [`Adam`](crate::neural_network::optimizers::Adam)
    ///
    /// # Returns
    ///
    /// - `Result<Self, Error>` - a new AdamW optimizer instance or an error
    ///
    /// # Notes
    ///
    /// Gradient clipping is disabled by default. Enable it with [`AdamW::with_global_clipnorm`]
    ///
    /// # Errors
    ///
    /// - `Error::InvalidParameter` - If the learning rate or epsilon is not positive and finite. Both beta values must be
    ///   finite in `[0, 1)`.
    ///   Weight decay must be nonnegative and finite.
    pub fn new(
        learning_rate: f32,
        beta1: f32,
        beta2: f32,
        epsilon: f32,
        weight_decay: f32,
    ) -> Result<Self, Error> {
        Ok(Self {
            core: AdamCore::new(learning_rate, beta1, beta2, epsilon, weight_decay, true)?,
        })
    }

    /// Enables a global gradient norm limit (disabled by default)
    ///
    /// The training loop scales gradients when their finite global L2 norm exceeds the
    /// threshold. A non-finite norm leaves gradients unchanged.
    ///
    /// # Parameters
    ///
    /// - `global_clipnorm` - clip-by-global-norm threshold. Must be positive and finite
    ///
    /// # Returns
    ///
    /// - `Result<Self, Error>` - the updated optimizer, or an error if `global_clipnorm` is not
    ///   positive and finite
    ///
    /// # Errors
    ///
    /// - `Error::InvalidParameter` - If the threshold is not positive and finite.
    pub fn with_global_clipnorm(self, global_clipnorm: f32) -> Result<Self, Error> {
        Ok(Self {
            core: self.core.with_global_clipnorm(global_clipnorm)?,
        })
    }
}

impl Optimizer for AdamW {
    fn global_clipnorm(&self) -> Option<f32> {
        self.core.global_clipnorm()
    }

    fn learning_rate(&self) -> f32 {
        self.core.learning_rate()
    }

    fn set_learning_rate(&mut self, learning_rate: f32) {
        self.core.set_learning_rate(learning_rate);
    }

    fn step(&mut self) {
        self.core.step();
    }

    fn update(
        &mut self,
        path: &LayerPath,
        layer: &mut dyn LayerBase,
        grads: &Grads,
        grad_scale: f32,
    ) {
        self.core.update(path, layer, grads, grad_scale);
    }
}