Struct OptimizationState 

pub struct OptimizationState<A: Float + 'static, D: Dimension + 'static> { /* private fields */ }

Current optimization state: parameters plus everything needed to move them.

Implementations

impl<A, D> OptimizationState<A, D>

where A: Float + ScalarOperand + Debug + Send + Sync + 'static, D: Dimension + 'static,

pub fn new( parameters: Array<A, D>, kind: HardwareOptimizerKind, base_learning_rate: A, accumulation_steps: usize, ) -> Self

Build a state around parameters, using the optimizer family kind.

accumulation_steps is clamped to at least 1: zero micro-batches per update would mean the parameters never move.

pub fn step(&mut self, gradients: &Array<A, D>) -> Result<HardwareStepReport<A>>

Apply one gradient.

With accumulation_steps == 1 this performs an optimizer update immediately. Otherwise the gradient is summed into the accumulator and the update happens once the window is full, using the mean of the window so the effective learning rate does not scale with the number of micro-batches.

A gradient whose shape differs from the parameters is reported rather than silently zipped against a truncated view.

pub fn parameters(&self) -> &Array<A, D>

Current parameters.

pub fn step_count(&self) -> usize

Number of optimizer updates applied so far.

pub fn optimizer_kind(&self) -> HardwareOptimizerKind

Optimizer family currently in use.

pub fn learning_rate(&self) -> A

Learning rate the next update will use.

pub fn base_learning_rate(&self) -> A

Rate this state was constructed with, for comparison against the live OptimizationState::learning_rate.

pub fn accumulation_steps(&self) -> usize

Micro-batches per optimizer update.

pub fn accumulated_micro_steps(&self) -> usize

Micro-batches currently held in the accumulator.

pub fn set_accumulation_steps(&mut self, accumulation_steps: usize)

Change the accumulation window.

Any partially accumulated window is discarded: mixing gradients averaged over different window sizes would silently change the effective learning rate of the next update.

pub fn set_lr_scheduler( &mut self, schedule: Box<dyn LearningRateScheduler<A> + Send + Sync>, )

Install a learning rate schedule.

The schedule owns the rate from this point on: every update reads LearningRateScheduler::get_learning_rate and pushes it into the optimizer, so a rate set directly on the optimizer would be overwritten.

base_learning_rate is deliberately left alone: it records what this state was constructed with, which is the only thing a later comparison against the live rate can be meaningful against.

pub fn rebuild_optimizer(&mut self, kind: HardwareOptimizerKind)

Replace the optimizer with a freshly built one of family kind.

This discards the accumulated moment state, which is why super::HardwareAwareOptimizer::optimize_for_hardware only does it before the first update.

pub fn set_optimizer( &mut self, kind: HardwareOptimizerKind, optimizer: Box<dyn Optimizer<A, D> + Send + Sync>, )

Install a caller-supplied optimizer, for the cases the built-in recommendation does not cover.

kind is what the state will report from OptimizationState::optimizer_kind; pass the family optimizer actually belongs to so the reported state stays truthful.

Trait Implementations

impl<A, D> Debug for OptimizationState<A, D>

where A: Float + Debug + 'static, D: Dimension + 'static,

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.