Skip to main content

DynOptimizer

burn_optim

Trait DynOptimizer 

Source 
pub trait DynOptimizer: Send + Sync {
    // Required methods
    fn step_dyn(
        &self,
        rank: usize,
        lr: LearningRate,
        tensor: BridgeTensor,
        grad: BridgeTensor,
        state: Option<DynState>,
    ) -> (BridgeTensor, Option<DynState>);
    fn to_device_dyn(&self, state: DynState, device: &Device) -> DynState;
    fn state_flatten(&self, prefix: &str, state: &DynState, out: &mut StateSink);
    fn state_unflatten(
        &self,
        rank: usize,
        prefix: &str,
        src: &mut StateSource,
        device: &Device,
    ) -> Option<DynState>;
}
Expand description

Object-safe view over an Optimizer, allowing ModuleOptimizer to stay non-generic. Rank-generic operations are dispatched on a runtime rank.

Required Methods§

Source

fn step_dyn( &self, rank: usize, lr: LearningRate, tensor: BridgeTensor, grad: BridgeTensor, state: Option<DynState>, ) -> (BridgeTensor, Option<DynState>)

Perform an optimizer step for a single parameter of the given rank.

Source

fn to_device_dyn(&self, state: DynState, device: &Device) -> DynState

Move a state to the given device.

Source

fn state_flatten(&self, prefix: &str, state: &DynState, out: &mut StateSink)

Decompose a state into named tensors and scalars under prefix.

Source

fn state_unflatten( &self, rank: usize, prefix: &str, src: &mut StateSource, device: &Device, ) -> Option<DynState>

Rebuild a state of the given rank from named tensors and scalars under prefix.

Returns None when the record does not contain a reconstructable state for this parameter (e.g. a truncated or foreign file); the caller leaves that parameter without state, so it is re-initialized on the next step.

Dyn Compatibility§

This trait is dyn compatible.

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

Implementors§