如何使用PyTorch实现自动混合精度训练？

本文共计1415个文字，预计阅读时间需要6分钟。

目录

1.torch.cuda.amp混合精度训练

2.Autocasting

2.1 torch.autocast 2.2 torch.cuda.amp.autocast

3.Gradient Scaling

3.1 使用案例：torch.cuda.amp混合精度训练

目录

• ​​1 torch.cuda.amp混合精度训练​​
• ​​2 Autocasting​​

• ​​2.1 torch.autocast​​
• ​​2.2 torch.cuda.amp.autocast​​

• ​​3 Gradient Scaling​​

• ​​3.1 使用案例​​

1 torch.cuda.amp混合精度训练

​ 混合精度训练提供了自适应的float32(单精度)与float16(半精度)数据适配，我们必须同时使用 ​​​torch.autocast​​​​ and ​​​torch.cuda.amp.GradScaler​​​​ 才能起到作用。然而，​​​torch.autocast​​​​和​​​GradScaler​​​​是模块化的，如果需要可以单独使用。