dfdx 0.13.0

Ergonomic auto differentiation in Rust, with pytorch like apis.
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

mod cpu_kernel;

#[cfg(feature = "cuda")]
mod cuda_kernel;

use crate::{
    shapes::{Dtype, Shape},
    tensor::*,
};

use super::WeightDecay;

/// Configuration of hyperparameters for [crate::optim::RMSprop].
#[derive(Debug, Clone, Copy)]
pub struct RMSpropConfig {
    /// Learning rate. Defaults to `1e-2`.
    pub lr: f64,

    /// Value for exponential moving average. Defaults to `0.9`.
    pub alpha: f64,

    /// Epsilon for stability. Defaults to `1e-8`.
    pub eps: f64,

    /// Optional momentum. Defaults to `None`.
    pub momentum: Option<f64>,

    /// Whether the avg should be centered by the grad's avg value.
    /// Defaults to `false`.
    pub centered: bool,

    /// Optional weight decay. Defaults to `None`.
    pub weight_decay: Option<WeightDecay>,
}

impl Default for RMSpropConfig {
    fn default() -> Self {
        Self {
            lr: 1e-2,
            alpha: 0.9,
            eps: 1e-8,
            momentum: None,
            centered: false,
            weight_decay: None,
        }
    }
}

pub trait RMSpropKernel<E: Dtype>: Storage<E> {
    fn rmsprop_kernel(
        &self,
        cfg: &RMSpropConfig,
        param: &mut Self::Vec,
        momentum: &mut Self::Vec,
        square_avg: &mut Self::Vec,
        grad_avg: &mut Self::Vec,
        grad: &Self::Vec,
    ) -> Result<(), Self::Err>;
}

impl RMSpropConfig {
    /// Update a single tensor using RMSprop.
    pub fn try_update<S: Shape, E: Dtype, D: RMSpropKernel<E>>(
        &self,
        param: &mut Tensor<S, E, D>,
        momentum: &mut D::Vec,
        square_avg: &mut D::Vec,
        grad_avg: &mut D::Vec,
        grad: &D::Vec,
    ) -> Result<(), D::Err> {
        param.device.rmsprop_kernel(
            self,
            std::sync::Arc::make_mut(&mut param.data),
            momentum,
            square_avg,
            grad_avg,
            grad,
        )
    }
}