黄梦迪，切坦·特库尔，迈克尔·卡里利

大多数深度学习框架，包括 PyTorch，默认使用 32 位浮点数（FP32）进行训练。然而，这并非实现许多深度学习模型完全精度的必要条件。2017 年，NVIDIA 研究人员开发了一种混合精度训练方法，该方法在训练网络时结合了单精度（FP32）和半精度（例如 FP16）格式，在相同的超参数下实现了与 FP32 训练相同的精度，并在 NVIDIA GPU 上获得了额外的性能优势：

短的训练时间；
更低的内存需求，使得能够使用更大的批量大小、更大的模型或更大的输入。

为了简化研究人员和从业者进行混合精度训练的用户体验，NVIDIA 于 2018 年开发了 Apex，这是一个轻量级的 PyTorch 扩展，具有自动混合精度（AMP）功能。该功能能够自动将某些 GPU 操作从 FP32 精度转换为混合精度，从而在保持准确性的同时提高性能。

在 PyTorch 1.6 版本发布时，NVIDIA 和 Facebook 的开发者将混合精度功能移至 PyTorch 核心作为 AMP 包 torch.cuda.amp。与 apex.amp 相比， torch.cuda.amp 更加灵活直观。以下是一些 apex.amp 已知的痛点， torch.cuda.amp 已经能够解决的：

保证 PyTorch 版本兼容性，因为它本是 PyTorch 的一部分
无需构建扩展
Windows 支持
检查点位运算精确保存/恢复
数据并行和进程内模型并行（尽管我们仍然推荐使用每个进程一个 GPU 的 torch.nn.DistributedDataParallel 作为最高效的方法）
梯度惩罚（双重反向）
torch.cuda.amp.autocast() 在未启用该功能的区域没有效果，因此它应该能够轻松处理以前在 apex.amp.initialize()（包括交叉验证）中遇到多个调用的案例。同一脚本中的多次收敛运行应各自使用一个新的 GradScaler 实例，但由于 GradScalers 轻量级且自包含，所以这并不是问题。
稀疏梯度支持

随着 AMP 被添加到 PyTorch 核心，我们已经开始弃用过程。我们已经将 apex.amp. 迁移到维护模式，并将支持使用 apex.amp. 的客户。然而，我们强烈建议 apex.amp 客户过渡到使用 torch.cuda.amp 从 PyTorch Core。

示例教程

请参阅官方文档了解使用方法：

例如：

import torch
# Creates once at the beginning of training
scaler = torch.cuda.amp.GradScaler()

for data, label in data_iter:
   optimizer.zero_grad()
   # Casts operations to mixed precision
   with torch.cuda.amp.autocast():
      loss = model(data)

   # Scales the loss, and calls backward()
   # to create scaled gradients
   scaler.scale(loss).backward()

   # Unscales gradients and calls
   # or skips optimizer.step()
   scaler.step(optimizer)

   # Updates the scale for next iteration
   scaler.update()

性能基准

在本节中，我们讨论了使用 AMP 在最新 NVIDIA GPU A100 和上一代 V100 GPU 上进行的混合精度训练的准确性和性能。混合精度性能与在 NVIDIA pytorch:20.06-py3 容器中运行深度学习工作负载时的 FP32 性能进行了比较。

准确性：AMP（FP16）、FP32

使用 AMP 进行深度学习训练的优势在于，模型收敛到相似的最终精度，同时提供改进的训练性能。为了说明这一点，对于 Resnet 50 v1.5 的训练，我们看到了以下精度结果，其中越高越好。请注意，以下精度数字是样本数字，可能存在 0.4%的运行到运行的变化。BERT、Transformer、ResNeXt-101、Mask-RCNN、DLRM 等其他模型的精度数字可以在 NVIDIA 深度学习示例 Github 上找到。

训练精度：NVIDIA DGX A100（8x A100 40GB）

轮次	混合精度 Top 1(%)	TF32 Top1(%)
90	76.93	76.85

训练准确率：NVIDIA DGX-1（8x V100 16GB）

轮次	混合精度 Top 1(%)	FP32 Top1(%)
50	76.25	76.26
90	77.09	77.01
250	78.42	78.30

加速性能：

在 NVIDIA V100 上使用 FP16 与在 V100 上使用 FP32

使用 FP16 的 AMP 是 V100 上深度学习训练中最具性能的选项。在表 1 中，我们可以观察到，对于各种模型，V100 上的 AMP 相对于 V100 上的 FP32 提供了 1.5 倍到 5.5 倍的加速，同时达到相同的最终精度。

图 2.在 NVIDIA 8xV100 上混合精度训练与 8xV100 GPU 上 FP32 训练的性能对比。柱状图表示 V100 AMP 相对于 V100 FP32 的速度提升因子。数值越高越好。

NVIDIA A100 上的 FP16 与 NVIDIA V100 上的 FP16

使用 FP16 的 AMP 仍然是 A100 上深度学习训练中最有效的选项。在图 3 中，我们可以观察到，对于各种模型，A100 上的 AMP 相对于 V100 上的 AMP 提供了 1.3 倍到 2.5 倍的速度提升，同时达到相同的最终精度。

图 3.在 NVIDIA 8xA100 上混合精度训练与 8xV100 GPU 上的性能对比。柱状图表示 A100 相对于 V100 的速度提升因子。数值越高越好。

行动号召

AMP 为 Nvidia Tensor Core GPU 上的深度学习训练工作负载提供了健康加速，尤其是在最新的 Ampere 代 A100 GPU 上。您可以从 NVIDIA 深度学习示例中开始使用 A100、V100、T4 和其他 GPU 进行 AMP 启用模型和模型脚本的实验。从 PyTorch NGC 容器版本 20.06 开始，NVIDIA PyTorch 支持原生 AMP。我们强烈建议现有 apex.amp 客户过渡到使用 torch.cuda.amp ，这是在最新 PyTorch 1.6 版本中可用的 PyTorch Core。

介绍原生 PyTorch 自动混合精度，在 NVIDIA GPU 上实现更快的训练