Skip to main content

Optimizer

tensorlogic_train

Trait Optimizer 

Source 
pub trait Optimizer {
    // Required methods
    fn step(
        &mut self,
        parameters: &mut HashMap<String, Array<f64, Ix2>>,
        gradients: &HashMap<String, Array<f64, Ix2>>,
    ) -> TrainResult<()>;
    fn zero_grad(&mut self);
    fn get_lr(&self) -> f64;
    fn set_lr(&mut self, lr: f64);
    fn state_dict(&self) -> HashMap<String, Vec<f64>>;
    fn load_state_dict(&mut self, state: HashMap<String, Vec<f64>>);
}
Expand description

Trait for optimizers.

Required Methods§

Source

fn step( &mut self, parameters: &mut HashMap<String, Array<f64, Ix2>>, gradients: &HashMap<String, Array<f64, Ix2>>, ) -> TrainResult<()>

Update parameters with computed gradients.

Source

fn zero_grad(&mut self)

Zero all gradients.

Source

fn get_lr(&self) -> f64

Get current learning rate.

Source

fn set_lr(&mut self, lr: f64)

Set learning rate.

Source

fn state_dict(&self) -> HashMap<String, Vec<f64>>

Get optimizer state for checkpointing.

Source

fn load_state_dict(&mut self, state: HashMap<String, Vec<f64>>)

Load optimizer state from checkpoint.

Implementors§