Struct MixedPrecisionTrainer 

pub struct MixedPrecisionTrainer { /* private fields */ }

Mixed precision trainer

Implementations

impl MixedPrecisionTrainer

pub fn new(config: MixedPrecisionConfig) -> Self

Create new mixed precision trainer

pub fn enable() -> Self

Enable mixed precision training

pub fn to_mixed_precision( &self, array: &Array2<Float>, operation_name: &str, ) -> MixedPrecisionArray

Convert array to mixed precision format

pub fn to_full_precision(&self, array: &MixedPrecisionArray) -> Array2<Float>

Convert mixed precision array back to FP32

pub fn scale_gradients(&self, gradients: &mut Array2<Float>)

Scale gradients to prevent underflow

pub fn unscale_gradients(&self, gradients: &mut Array2<Float>) -> bool

Unscale gradients after backward pass

pub fn update_scale(&mut self, overflow_detected: bool)

Update loss scale based on overflow detection

pub fn should_skip_step(&self) -> bool

Check if current step should be skipped due to overflow

pub fn get_loss_scale(&self) -> Float

Get current loss scale

pub fn train_ensemble_mixed_precision<F>( &mut self, x: &Array2<Float>, y: &Array1<Int>, n_estimators: usize, train_fn: F, ) -> Result<Vec<Array1<Float>>>

where F: FnMut(&MixedPrecisionArray, &Array1<Int>) -> Result<Array1<Float>>,

Train ensemble with mixed precision

pub fn scaler_state(&self) -> &ScalerState

Get scaler state

pub fn reset_scaler(&mut self)

Reset scaler state