Struct MixedPrecisionTrainer 

pub struct MixedPrecisionTrainer { /* private fields */ }

Mixed precision training manager.

Implementations

impl MixedPrecisionTrainer

pub fn new(mode: PrecisionMode, scaler: LossScaler) -> Self

Creates a new mixed precision trainer.

pub fn register_weights(&mut self, name: String, weights: Array2<f32>)

Registers weights to maintain master copy.

pub fn cast_to_working_precision(&self, weights: &Array2<f32>) -> Array2<f32>

Converts FP32 weights to working precision.

pub fn scale_loss(&mut self, loss: f32) -> f32

Scales loss for backward pass.

pub fn unscale_and_check_gradients( &mut self, gradients: &mut HashMap<String, Array2<f32>>, ) -> TrainResult<(bool, bool)>

Unscales and checks gradients for overflow.

Returns

(should_step, overflow_detected)

pub fn update_master_weights(&mut self, updates: &HashMap<String, Array2<f32>>)

Updates master weights from working precision weights.

pub fn mode(&self) -> PrecisionMode

Gets the current precision mode.

pub fn current_scale(&self) -> f32

Gets the current loss scale.

pub fn stats(&self) -> &MixedPrecisionStats

Gets training statistics.

pub fn reset_stats(&mut self)

Resets statistics.