Trait Optimizer
pub trait Optimizer:
Send
+ Sync
+ Clone
+ 'static {
type State<const D: usize>: Clone + RecordState;
// Required methods
fn step<const D: usize>(
&self,
lr: f64,
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§
type State<const D: usize>: Clone + RecordState
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§
fn step<const D: usize>(
&self,
lr: f64,
tensor: Tensor<D>,
grad: Tensor<D>,
state: Option<Self::State<D>>,
) -> (Tensor<D>, Option<Self::State<D>>)
fn step<const D: usize>( &self, lr: f64, 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.
Dyn Compatibility§
This trait is not dyn compatible.
In older versions of Rust, dyn compatibility was called "object safety".