NativeMixedPrecisionPlugin¶
- class pytorch_lightning.plugins.precision.NativeMixedPrecisionPlugin(precision, device, scaler=None)[source]¶
Bases:
pytorch_lightning.plugins.precision.mixed.MixedPrecisionPlugin
Plugin for Native Mixed Precision (AMP) training with
torch.autocast
.- Parameters
precision¶ (
Union
[str
,int
]) – Whether to usetorch.float16
(16
) ortorch.bfloat16
('bf16'
).scaler¶ (
Optional
[GradScaler
]) – An optionaltorch.cuda.amp.GradScaler
to use.
- on_load_checkpoint(checkpoint)[source]¶
Called by Lightning to restore your model. If you saved something with
on_save_checkpoint()
this is your chance to restore this.Example:
def on_load_checkpoint(self, checkpoint): # 99% of the time you don't need to implement this method self.something_cool_i_want_to_save = checkpoint['something_cool_i_want_to_save']
Note
Lightning auto-restores global step, epoch, and train state including amp scaling. There is no need for you to restore anything regarding training.
- on_save_checkpoint(checkpoint)[source]¶
Called by Lightning when saving a checkpoint to give you a chance to store anything else you might want to save.
- Parameters
checkpoint¶ (
Dict
[str
,Any
]) – The full checkpoint dictionary before it gets dumped to a file. Implementations of this hook can insert additional data into this dictionary.- Return type
Example:
def on_save_checkpoint(self, checkpoint): # 99% of use cases you don't need to implement this method checkpoint['something_cool_i_want_to_save'] = my_cool_pickable_object
Note
Lightning saves all aspects of training (epoch, global step, etc…) including amp scaling. There is no need for you to store anything about training.