Struct MixedPrecisionManager

pub struct MixedPrecisionManager<F: Float + Debug + ScalarOperand + FromPrimitive + Send + Sync + Display, LF: Float + Debug + ScalarOperand + Send + Sync> {
    pub config: MixedPrecisionConfig,
    /* private fields */
}

Mixed precision training manager

Fields

config: MixedPrecisionConfig

Configuration for mixed precision training

Implementations

impl<F: Float + Debug + ScalarOperand + FromPrimitive + Send + Sync + Display, LF: Float + Debug + ScalarOperand + Send + Sync> MixedPrecisionManager<F, LF>

pub fn new(config: MixedPrecisionConfig) -> Self

Create a new mixed precision manager

pub fn initialize<L: Layer<F> + Clone + Send + Sync + 'static>( &mut self, model: &L, ) -> Result<()>

Initialize the manager with a model

pub fn to_low_precision<L: Layer<F> + ?Sized, L2: Layer<LF> + ?Sized>( &self, high_precision_model: &L, low_precision_model: &mut L2, ) -> Result<()>

Convert a model’s parameters to low precision

pub fn gradients_to_high_precision( &mut self, low_precision_gradients: &[Array<LF, IxDyn>], ) -> Result<Vec<Array<F, IxDyn>>>

Convert gradients from low to high precision

pub fn train_step<LM: Layer<LF> + ?Sized, HM: Layer<F> + ParamLayer<F> + ?Sized, O: Optimizer<F> + OptimizerStep<F>>( &mut self, low_precision_model: &mut LM, high_precision_model: &mut HM, optimizer: &mut O, inputs: &Array<F, IxDyn>, targets: &Array<F, IxDyn>, loss_fn: &dyn Loss<F>, ) -> Result<F>

Train a step with mixed precision

pub fn get_loss_scale(&self) -> F

Get the current loss scale

pub fn get_success_counter(&self) -> usize

Get the number of consecutive successful steps

pub fn overflow_occurred(&self) -> bool

Check if overflow occurred in the last step

pub fn train_epoch<LM: Layer<LF> + ?Sized, HM: ParamLayer<F> + ?Sized, O: Optimizer<F>, D: Dataset<F> + Clone>( &mut self, low_precision_model: &mut LM, high_precision_model: &mut HM, optimizer: &mut O, dataset: &D, loss_fn: &dyn Loss<F>, batch_size: usize, shuffle: bool, ) -> Result<(F, usize)>

Train for one epoch with mixed precision