Trait UnifiedOptimizer

pub trait UnifiedOptimizer<A: Float> {
    // Required methods
    fn config(&self) -> &OptimizerConfig<A>;
    fn step_param<D: Dimension>(
        &mut self,
        param: &mut Parameter<A, D>,
    ) -> Result<()>
       where A: ScalarOperand + Debug;
    fn set_lr(&mut self, lr: A);
    fn get_lr(&self) -> A;
    fn state_dict(&self) -> HashMap<String, Vec<u8>>;
    fn load_state_dict(
        &mut self,
        state_dict: HashMap<String, Vec<u8>>,
    ) -> Result<()>;

    // Provided methods
    fn step_params<D: Dimension>(
        &mut self,
        params: &mut [Parameter<A, D>],
    ) -> Result<()>
       where A: ScalarOperand + Debug { ... }
    fn zero_grad<D: Dimension>(&self, params: &mut [Parameter<A, D>]) { ... }
}

Unified optimizer interface

Required Methods

fn config(&self) -> &OptimizerConfig<A>

Get optimizer configuration

fn step_param<D: Dimension>( &mut self, param: &mut Parameter<A, D>, ) -> Result<()>

where A: ScalarOperand + Debug,

Update a single parameter

fn set_lr(&mut self, lr: A)

Update learning rate

fn get_lr(&self) -> A

Get current learning rate

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

State dictionary for serialization

fn load_state_dict( &mut self, state_dict: HashMap<String, Vec<u8>>, ) -> Result<()>

Load state from dictionary

Provided Methods

fn step_params<D: Dimension>( &mut self, params: &mut [Parameter<A, D>], ) -> Result<()>

where A: ScalarOperand + Debug,

Update multiple parameters

fn zero_grad<D: Dimension>(&self, params: &mut [Parameter<A, D>])

Zero gradients for all parameters

Dyn Compatibility

This trait is not dyn compatible.

In older versions of Rust, dyn compatibility was called "object safety", so this trait is not object safe.

Implementors

impl<A: Float + ScalarOperand + Debug> UnifiedOptimizer<A> for UnifiedAdam<A>

impl<A: Float + ScalarOperand + Debug> UnifiedOptimizer<A> for UnifiedSGD<A>