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
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130

//! Stochastic gradient descent.

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

use super::state::place;
use super::{Optimizer, OptimizerState};

/// SGD with optional momentum, Nesterov, and weight decay.
///
/// ```
/// # use fastnn::prelude::*;
/// # let model = Sequential::new().add(Linear::new(4, 2));
/// let mut opt = SGD::new(model.parameters(), 0.1)
///     .momentum(0.9)
///     .weight_decay(5e-4);
/// ```
pub struct SGD {
    params: Vec<Param>,
    velocity: Vec<Option<Tensor>>,
    lr: f32,
    momentum: f32,
    dampening: f32,
    weight_decay: f32,
    nesterov: bool,
    steps: u64,
}

impl SGD {
    pub fn new(params: Vec<Param>, lr: f32) -> SGD {
        SGD {
            velocity: vec![None; params.len()],
            params,
            lr,
            momentum: 0.0,
            dampening: 0.0,
            weight_decay: 0.0,
            nesterov: false,
            steps: 0,
        }
    }

    /// Fraction of the previous update carried forward.
    pub fn momentum(mut self, momentum: f32) -> SGD {
        self.momentum = momentum;
        self
    }

    /// Damping on each new gradient entering the velocity.
    pub fn dampening(mut self, dampening: f32) -> SGD {
        self.dampening = dampening;
        self
    }

    /// L2 penalty added to the gradient.
    pub fn weight_decay(mut self, decay: f32) -> SGD {
        self.weight_decay = decay;
        self
    }

    /// Use Nesterov momentum.
    pub fn nesterov(mut self, nesterov: bool) -> SGD {
        self.nesterov = nesterov;
        self
    }
}

impl Optimizer for SGD {
    fn step(&mut self) {
        self.steps += 1;
        for (index, param) in self.params.iter().enumerate() {
            if !param.is_trainable() {
                continue;
            }
            let Some(grad) = param.grad() else { continue };
            let value = param.value();

            let grad = if self.weight_decay == 0.0 {
                grad
            } else {
                grad.add(&value.mul_scalar(self.weight_decay))
            };

            let update = if self.momentum == 0.0 {
                grad
            } else {
                let velocity = match &self.velocity[index] {
                    Some(previous) => previous
                        .mul_scalar(self.momentum)
                        .add(&grad.mul_scalar(1.0 - self.dampening)),
                    None => grad.clone(),
                };
                let update = if self.nesterov {
                    grad.add(&velocity.mul_scalar(self.momentum))
                } else {
                    velocity.clone()
                };
                self.velocity[index] = Some(velocity);
                update
            };

            param.set_value(value.sub(&update.mul_scalar(self.lr)));
        }
    }

    fn parameters(&self) -> &[Param] {
        &self.params
    }

    fn lr(&self) -> f32 {
        self.lr
    }

    fn set_lr(&mut self, lr: f32) {
        self.lr = lr;
    }

    fn state(&self) -> OptimizerState {
        let mut state = OptimizerState::new(self.steps);
        state.put("velocity", &self.velocity);
        state
    }

    fn load_state(&mut self, state: OptimizerState) -> Result<()> {
        self.velocity = place(state.take("velocity", self.params.len())?, &self.params);
        self.steps = state.steps;
        Ok(())
    }
}