syntaxdot 0.2.0

Neural sequence labeler
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

use tch::nn::{self};
use tch::Tensor;

mod grad;
pub use grad::ZeroGrad;

mod grad_scale;
pub use grad_scale::GradScaler;

pub trait Optimizer {
    /// Perform a backward step on the given loss.
    fn backward_step(&mut self, loss: &Tensor);

    /// Set the learning rate for a parameter group.
    fn set_lr_group(&mut self, group: usize, learning_rate: f64);

    /// Set the weight decay for a parameter group.
    fn set_weight_decay_group(&mut self, group: usize, weight_decay: f64);

    /// Perform an update step.
    ///
    /// It is generally recommended to use `backward_step`, since it
    /// computes the gradients and performs any loss scaling (if
    /// necessary).
    fn step(&mut self);

    /// Get the trainable variables associated with the optimizer.
    fn trainable_variables(&self) -> Vec<Tensor>;
}

impl<C> Optimizer for nn::Optimizer<C> {
    fn backward_step(&mut self, loss: &Tensor) {
        nn::Optimizer::backward_step(self, loss)
    }

    fn set_lr_group(&mut self, group: usize, learning_rate: f64) {
        nn::Optimizer::set_lr_group(self, group, learning_rate)
    }

    fn set_weight_decay_group(&mut self, group: usize, weight_decay: f64) {
        nn::Optimizer::set_weight_decay_group(self, group, weight_decay)
    }

    fn step(&mut self) {
        nn::Optimizer::step(self)
    }

    fn trainable_variables(&self) -> Vec<Tensor> {
        nn::Optimizer::trainable_variables(self)
    }
}