MixedPrecisionPlugin¶
- class lightning.pytorch.plugins.precision.MixedPrecisionPlugin(precision, device, scaler=None)[source]¶
Bases:
lightning.pytorch.plugins.precision.precision_plugin.PrecisionPlugin
Plugin for Automatic Mixed Precision (AMP) training with
torch.autocast
.- Parameters
precision¶ (
Literal
[‘16-mixed’, ‘bf16-mixed’]) – Whether to usetorch.float16
(16
) ortorch.bfloat16
('bf16'
).scaler¶ (
Optional
[GradScaler
]) – An optionaltorch.cuda.amp.GradScaler
to use.
- clip_gradients(optimizer, clip_val=0.0, gradient_clip_algorithm=GradClipAlgorithmType.NORM)[source]¶
Clips the gradients.
- Return type
- load_state_dict(state_dict)[source]¶
Called when loading a checkpoint, implement to reload precision plugin state given precision plugin state_dict.
- optimizer_step(optimizer, model, closure, **kwargs)[source]¶
Hook to run the optimizer step.
- Return type