hodu_nn 0.2.1

hodu nn
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

use crate::optimizer::Optimizer;
use hodu_core::{error::HoduResult, scalar::Scalar, tensor::Tensor};

#[derive(Optimizer, Clone)]
pub struct SGD {
    learning_rate: Scalar,
}

impl SGD {
    pub fn new(learning_rate: impl Into<Scalar>) -> Self {
        Self {
            learning_rate: learning_rate.into(),
        }
    }

    pub fn step(&mut self, parameters: &mut [&mut Tensor]) -> HoduResult<()> {
        for param in parameters.iter_mut() {
            let grad = param.grad()?;
            let lr = self.learning_rate.to_dtype(grad.dtype());

            param.set_(&param.sub(&grad.mul_scalar(lr)?)?)?;
        }
        Ok(())
    }

    pub fn set_learning_rate(&mut self, learning_rate: impl Into<Scalar>) {
        self.learning_rate = learning_rate.into();
    }
}