Struct MixedPrecisionOptimizer 

pub struct MixedPrecisionOptimizer<O: Optimizer> { /* private fields */ }

Mixed precision optimizer wrapper

Implementations

impl<O: Optimizer> MixedPrecisionOptimizer<O>

pub fn new(optimizer: O, config: MixedPrecisionConfig) -> Self

Create a new mixed precision optimizer wrapper

pub fn with_defaults(optimizer: O) -> Self

Create a mixed precision optimizer with default configuration

pub fn get_loss_scale(&self) -> f32

Get the current loss scale

pub fn is_enabled(&self) -> bool

Check if mixed precision is enabled

pub fn initialize_master_weights( &mut self, params: &[Arc<RwLock<Tensor>>], ) -> Result<()>

Initialize master weights for mixed precision training

pub fn scale_loss(&mut self, loss: &mut Tensor) -> Result<()>

Scale loss for backward pass

pub fn unscale_gradients( &mut self, params: &[Arc<RwLock<Tensor>>], ) -> Result<bool>

Unscale gradients before optimizer step

pub fn update_master_weights( &mut self, params: &[Arc<RwLock<Tensor>>], ) -> Result<()>

Update master weights from fp16 parameters

pub fn copy_master_to_params( &mut self, params: &[Arc<RwLock<Tensor>>], ) -> Result<()>

Copy master weights back to fp16 parameters

pub fn inner(&self) -> &O

Get the underlying optimizer

pub fn inner_mut(&mut self) -> &mut O

Get the underlying optimizer mutably

Trait Implementations

impl<O: Optimizer> Optimizer for MixedPrecisionOptimizer<O>

fn step(&mut self) -> OptimizerResult<()>

Perform a single optimization step

fn zero_grad(&mut self)

Zero all gradients

fn get_lr(&self) -> Vec<f32>

Get the current learning rate

fn set_lr(&mut self, lr: f32)

Set one learning rate for every parameter group (broadcast).

fn add_param_group( &mut self, params: Vec<Arc<RwLock<Tensor>>>, options: HashMap<String, f32>, )

Add a parameter group

fn parameters(&self) -> Vec<Arc<RwLock<Tensor>>>

Get the parameter tensors managed by this optimizer.

fn state_dict(&self) -> OptimizerResult<OptimizerState>

Get state dict for serialization

fn load_state_dict(&mut self, state: OptimizerState) -> OptimizerResult<()>

Load state dict

fn set_lrs(&mut self, lrs: &[f32])

Set the learning rate of each parameter group individually.