burn 0.3.0

BURN: Burn Unstoppable Rusty Neurons
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

use crate as burn;

use super::{load_state_gradients, register_state_gradients};
use crate::config::Config;
use crate::module::{ParamId, StateNamed};
use crate::tensor::backend::ADBackend;
use crate::tensor::{ElementConversion, Tensor};
use burn_tensor::backend::Gradients;

/// Configuration to create [WeightDecay](WeightDecay).
#[derive(Config)]
pub struct WeightDecayConfig {
    /// L2 penalty.
    pub penalty: f64,
}

/// Weight decay implementation that transforms gradients.
pub struct WeightDecay<B: ADBackend> {
    penalty: B::Elem,
    gradients: B::Gradients,
}

impl<B: ADBackend> WeightDecay<B> {
    pub fn new(config: &WeightDecayConfig) -> Self {
        Self {
            penalty: config.penalty.to_elem(),
            gradients: B::Gradients::empty(),
        }
    }

    pub fn transform<const D: usize>(
        &mut self,
        id: &ParamId,
        grad: Tensor<B::InnerBackend, D>,
    ) -> Tensor<B::InnerBackend, D> {
        let id = id.to_string();

        let grad = match self.gradients.get::<D>(&id) {
            Some(grad_last_step) => grad_last_step.mul_scalar(self.penalty).add(&grad),
            None => grad,
        };

        // Update gradients
        self.gradients.register(id, grad.clone());

        grad
    }
    pub fn register_state<const D: usize>(&self, id: &ParamId, state: &mut StateNamed<B::Elem>) {
        register_state_gradients::<D, B, _>(id, state, &self.gradients, Self::state_key);
    }

    pub fn load_state<const D: usize>(
        &mut self,
        id: &ParamId,
        state: &StateNamed<B::Elem>,
        device: &B::Device,
    ) {
        load_state_gradients::<D, B, _>(id, state, &mut self.gradients, Self::state_key, device);
    }

    fn state_key(id: &str) -> String {
        format!("weight-decay-{}", id)
    }
}