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
use burn_core as burn;
use burn::config::Config;
use burn::record::Record;
use burn::tensor::backend::Backend;
use burn::tensor::{ElementConversion, Tensor};
/// Configuration to create [momentum](Momentum).
#[derive(Config, Debug)]
pub struct MomentumConfig {
/// Momentum factor
#[config(default = 0.9)]
pub momentum: f64,
/// Dampening factor.
#[config(default = 0.1)]
pub dampening: f64,
/// Enables Nesterov momentum, see [On the importance of initialization and
/// momentum in deep learning](http://www.cs.toronto.edu/~hinton/absps/momentum.pdf).
#[config(default = false)]
pub nesterov: bool,
}
/// State of [momentum](Momentum).
#[derive(Record, Clone, new)]
pub struct MomentumState<B: Backend, const D: usize> {
velocity: Tensor<B, D>,
}
/// Momentum implementation that transforms gradients.
#[derive(Clone)]
pub struct Momentum<B: Backend> {
momentum: B::FloatElem,
dampening: f64,
nesterov: bool,
}
impl<B: Backend> Momentum<B> {
/// Creates a new [momentum](Momentum) from a [config](MomentumConfig).
pub fn new(config: &MomentumConfig) -> Self {
Self {
momentum: config.momentum.elem(),
dampening: config.dampening,
nesterov: config.nesterov,
}
}
/// Transforms a gradient.
///
/// # Arguments
///
/// * `grad` - Gradient to transform.
/// * `state` - State of the optimizer.
///
/// # Returns
///
/// * `grad` - Transformed gradient.
/// * `state` - State of the optimizer.
pub fn transform<const D: usize>(
&self,
grad: Tensor<B, D>,
state: Option<MomentumState<B, D>>,
) -> (Tensor<B, D>, MomentumState<B, D>) {
let velocity = if let Some(state) = state {
grad.clone()
.mul_scalar(1.0 - self.dampening)
.add(state.velocity.mul_scalar(self.momentum))
} else {
grad.clone()
};
let grad = match self.nesterov {
true => velocity.clone().mul_scalar(self.momentum).add(grad),
false => velocity.clone(),
};
(grad, MomentumState::new(velocity))
}
}
impl<B: Backend, const D: usize> MomentumState<B, 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: &B::Device) -> Self {
self.velocity = self.velocity.to_device(device);
self
}
}