fastnn 0.3.3

A robust deep learning library in Rust with CUDA kernels
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

//! Exponential moving average of the weights, for evaluating and shipping.

use crate::nn::Param;
use crate::tensor::Tensor;

/// Shadow weights: `shadow = d * shadow + (1 - d) * param` after every step.
///
/// ```
/// # use fastnn::prelude::*;
/// # use fastnn::optim::Ema;
/// # let model = Sequential::new().add(Linear::new(4, 2));
/// # let x = Tensor::randn(&[1, 4]);
/// let mut opt = SGD::new(model.parameters(), 0.1);
/// let mut ema = Ema::new(model.parameters(), 0.999);
///
/// // ... loss.backward(); opt.step(); ...
/// ema.update();
///
/// ema.swap();     // evaluate with the averaged weights
/// let logits = no_grad(|| model.forward(&x));
/// ema.swap();     // back to the raw weights; training continues
/// ```
pub struct Ema {
    params: Vec<Param>,
    shadow: Vec<Tensor>,
    decay: f32,
}

impl Ema {
    /// `decay` is usually between 0.99 and 0.9999.
    pub fn new(params: Vec<Param>, decay: f32) -> Ema {
        assert!(
            (0.0..1.0).contains(&decay),
            "decay must be in [0, 1), got {decay}"
        );
        // Seeded at the current weights, not zero, to avoid a long bias toward the origin.
        let shadow = params.iter().map(|p| p.value().detach()).collect();
        Ema {
            params,
            shadow,
            decay,
        }
    }

    /// Fold the current weights into the average. Call after `opt.step()`.
    pub fn update(&mut self) {
        for (shadow, param) in self.shadow.iter_mut().zip(&self.params) {
            let value = param.value();
            *shadow = shadow
                .mul_scalar(self.decay)
                .add(&value.mul_scalar(1.0 - self.decay));
        }
    }

    /// Swap the model's weights with the averaged ones. Call again to swap back.
    pub fn swap(&mut self) {
        for (shadow, param) in self.shadow.iter_mut().zip(&self.params) {
            let raw = param.value().detach();
            param.set_value(shadow.clone());
            *shadow = raw;
        }
    }

    /// The averaged tensors in parameter order.
    pub fn averaged(&self) -> &[Tensor] {
        &self.shadow
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn update_follows_the_definition_exactly() {
        let param = Param::new(Tensor::from_vec(vec![1.0], &[1]));
        let mut ema = Ema::new(vec![param.clone()], 0.9);

        param.set_value(Tensor::from_vec(vec![3.0], &[1]));
        ema.update();

        assert!((ema.averaged()[0].to_vec()[0] - 1.2).abs() < 1e-6);
    }

    #[test]
    fn double_swap_is_identity() {
        let param = Param::new(Tensor::from_vec(vec![5.0], &[1]));
        let mut ema = Ema::new(vec![param.clone()], 0.99);
        param.set_value(Tensor::from_vec(vec![7.0], &[1]));

        ema.swap();
        assert_eq!(
            param.value().to_vec()[0],
            5.0,
            "first swap shows the average"
        );
        ema.swap();
        assert_eq!(
            param.value().to_vec()[0],
            7.0,
            "second swap restores training weights"
        );
    }
}