burn_optim/optim/
decay.rs

1use burn_core as burn;
2
3use crate::RecordState;
4use burn::config::Config;
5use burn::tensor::Device;
6use burn::tensor::Tensor;
7
8/// Configuration to create [weight decay](WeightDecay).
9#[derive(Config, Debug)]
10pub struct WeightDecayConfig {
11    /// L2 penalty.
12    pub penalty: f32,
13}
14
15/// State of [weight decay](WeightDecay).
16#[derive(RecordState, Clone, new)]
17pub struct WeightDecayState<const D: usize> {
18    pub(crate) grad_last_step: Tensor<D>,
19}
20
21/// Weight decay implementation that transforms gradients.
22#[derive(Clone)]
23pub struct WeightDecay {
24    penalty: f32,
25}
26
27impl WeightDecay {
28    /// Creates a new [weight decay](WeightDecay) from a [config](WeightDecayConfig).
29    pub fn new(config: &WeightDecayConfig) -> Self {
30        Self {
31            penalty: config.penalty,
32        }
33    }
34
35    /// Transforms a gradient.
36    ///
37    /// # Arguments
38    ///
39    /// * `grad` - Gradient to transform.
40    /// * `tensor` - Tensor param of the last iteration.
41    ///
42    /// # Returns
43    ///
44    /// * `grad` - Transformed gradient.
45    pub fn transform<const D: usize>(&self, grad: Tensor<D>, tensor: Tensor<D>) -> Tensor<D> {
46        tensor.mul_scalar(self.penalty).add(grad)
47    }
48}
49
50impl<const D: usize> WeightDecayState<D> {
51    /// Moves the state to a device.
52    ///
53    /// # Arguments
54    ///
55    /// * `device` - Device to move the state to.
56    ///
57    /// # Returns
58    ///
59    /// * `self` - Moved state.
60    pub fn to_device(mut self, device: &Device) -> Self {
61        self.grad_last_step = self.grad_last_step.to_device(device);
62        self
63    }
64}
burn_optim/optim/decay.rs

burn_optim/optim/
decay.rs
