作者：Mengdi Huang、Chetan Tekur、Michael Carilli

大多数深度学习框架（包括 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。torch.cuda.amp 相比于 apex.amp 更加灵活和直观。torch.cuda.amp 能够解决 apex.amp 中一些已知的痛点：

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

随着 AMP 被添加到 PyTorch 核心，我们已开始弃用 apex.amp 的进程。我们已将 apex.amp 置于维护模式，并将继续为使用 apex.amp 的客户提供支持。然而，我们强烈鼓励 apex.amp 用户迁移到使用 PyTorch 核心中的 torch.cuda.amp。

示例演练

请参阅官方文档了解用法

示例

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()

性能基准测试

在本节中，我们将讨论在最新的 NVIDIA GPU A100 以及上一代 V100 GPU 上使用 AMP 进行混合精度训练的精度和性能。在 NGC 的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

加速性能

FP16 在 NVIDIA V100 上 vs. FP32 在 V100 上

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

图 2. 在 NVIDIA 8xV100 上进行混合精度训练与在 8xV100 GPU 上进行 FP32 训练的性能。条形图表示 V100 AMP 相较于 V100 FP32 的加速因子。数值越高越好。

FP16 在 NVIDIA A100 上 vs. FP16 在 V100 上

使用 FP16 的 AMP 仍然是 A100 上深度学习训练性能最佳的选项。在图 3 中，我们可以观察到，对于各种模型，AMP 在 A100 上的性能比 AMP 在 V100 上提高了 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 模型和模型脚本进行实验。支持原生 AMP 的 NVIDIA PyTorch 可从PyTorch NGC 容器版本 20.06 中获取。我们强烈鼓励现有的 apex.amp 用户迁移到使用最新PyTorch 1.6 版本中提供的 PyTorch Core 中的 torch.cuda.amp。

引入 PyTorch 原生自动混合精度，可在 NVIDIA GPU 上加速训练