Trait Optimizer 

pub trait Optimizer {
    // Required methods
    fn step(&mut self, params: &mut [Tensor]);
    fn lr(&self) -> f32;
    fn set_lr(&mut self, lr: f32);

    // Provided methods
    fn step_refs(&mut self, params: &mut [&mut Tensor]) { ... }
    fn zero_grad(&mut self, params: &mut [Tensor]) { ... }
    fn zero_grad_refs(&mut self, params: &mut [&mut Tensor]) { ... }
}

Trait for optimization algorithms

Required Methods

fn step(&mut self, params: &mut [Tensor])

Perform a single optimization step

fn lr(&self) -> f32

Get learning rate

fn set_lr(&mut self, lr: f32)

Set learning rate

Provided Methods

fn step_refs(&mut self, params: &mut [&mut Tensor])

Perform optimization step on referenced parameters

This is useful when parameters are borrowed from a model

fn zero_grad(&mut self, params: &mut [Tensor])

Zero out all gradients

fn zero_grad_refs(&mut self, params: &mut [&mut Tensor])

Zero gradients on referenced parameters

Implementors

impl Optimizer for Adam

impl Optimizer for AdamW

impl Optimizer for SGD