由浅入深的混合精度训练教程

文章预览

本文转载自 | PaperWeekly 作者 | serendipity 单位 | 同济大学 研究方向 | 行人搜索 2022 年的当下，混合精度 (Automatically Mixed Precision, AMP) 训练已经成为了炼丹师的标配工具，仅仅只需几行代码，就能让显存占用减半，训练速度加倍。  AMP 技术是由百度和 NIVDIA 团队在 2017 年提出的 (Mixed Precision Training [1])，该成果发表在 ICLR 上。PyTorch 1.6之前，大家都是用 NVIDIA 的 apex [2] 库来实现 AMP 训练。1.6 版本之后，PyTorch 出厂自带 AMP。 这篇文章由浅入深地讲解了: 如何在 PyTorch 中使用 AMP、AMP 的原理、AMP 的代码实现。 1 『如何在pytorch中使用AMP』 如果你是新手，只是想简单地试用一下 AMP，只需要将相关训练代码 output = net(input) loss = loss_fn(output, target) loss.backward() optimizer.step() optimizer.zero_grad() 修改如下即可。 with  torch.cuda.amp.autocast():     output = net(input)     l ………………………………