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
use burn_core as burn;
use crate::RecordState;
use burn::config::Config;
use burn::tensor::Device;
use burn::tensor::Tensor;
/// Configuration to create [weight decay](WeightDecay).
#[derive(Config, Debug)]
pub struct WeightDecayConfig {
/// L2 penalty.
pub penalty: f32,
}
/// State of [weight decay](WeightDecay).
#[derive(RecordState, Clone, new)]
pub struct WeightDecayState<const D: usize> {
pub(crate) grad_last_step: Tensor<D>,
}
/// Weight decay implementation that transforms gradients.
#[derive(Clone)]
pub struct WeightDecay {
penalty: f32,
}
impl WeightDecay {
/// Creates a new [weight decay](WeightDecay) from a [config](WeightDecayConfig).
pub fn new(config: &WeightDecayConfig) -> Self {
Self {
penalty: config.penalty,
}
}
/// Transforms a gradient.
///
/// # Arguments
///
/// * `grad` - Gradient to transform.
/// * `tensor` - Tensor param of the last iteration.
///
/// # Returns
///
/// * `grad` - Transformed gradient.
pub fn transform<const D: usize>(&self, grad: Tensor<D>, tensor: Tensor<D>) -> Tensor<D> {
tensor.mul_scalar(self.penalty).add(grad)
}
}
impl<const D: usize> WeightDecayState<D> {
/// Moves the state to a device.
///
/// # Arguments
///
/// * `device` - Device to move the state to.
///
/// # Returns
///
/// * `self` - Moved state.
pub fn to_device(mut self, device: &Device) -> Self {
self.grad_last_step = self.grad_last_step.to_device(device);
self
}
}