Skip to main content

Optimizer

burn_optim

Trait Optimizer 

Source 
pub trait Optimizer:
    Send
    + Sync
    + Clone
    + 'static {
    type State<const D: usize>: Clone + RecordState;

    // Required methods
    fn step<const D: usize>(
        &self,
        lr: LearningRate,
        tensor: Tensor<D>,
        grad: Tensor<D>,
        state: Option<Self::State<D>>,
    ) -> (Tensor<D>, Option<Self::State<D>>);
    fn to_device<const D: usize>(
        state: Self::State<D>,
        device: &Device,
    ) -> Self::State<D>;
}
Expand description

An opinionated trait to simplify the process of implementing an optimizer.

Implementations don’t have to handle missing gradients, loading and exporting records, navigate the module parameter structure, handle tracked and untracked tensors, and the likes. Wrap one in a ModuleOptimizer to optimize a whole module.

Required Associated Types§

Source

type State<const D: usize>: Clone + RecordState

The state of the optimizer for a single parameter of rank D.

It implements RecordState (which itself requires Send + Sync + 'static) so it can be decomposed into named tensors and scalars for the burnpack format.

Required Methods§

Source

fn step<const D: usize>( &self, lr: LearningRate, tensor: Tensor<D>, grad: Tensor<D>, state: Option<Self::State<D>>, ) -> (Tensor<D>, Option<Self::State<D>>)

The optimizer step is performed for one tensor at a time with its gradient and state.

Note that the state is passed as parameter, so implementations don’t have to handle the saving and loading of recorded states.

Source

fn to_device<const D: usize>( state: Self::State<D>, device: &Device, ) -> Self::State<D>

Change the device of the state.

This function will be called accordingly to have the state on the same device as the gradient and the tensor when the step function is called.

Dyn Compatibility§

This trait is not dyn compatible.

In older versions of Rust, dyn compatibility was called "object safety".

Implementors§