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

缩短训练时间；
降低内存需求，从而实现更大的批处理大小、更大的模型或更大的输入。

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

对于 PyTorch 1.6 版本，NVIDIA 和 Facebook 的开发人员将混合精度功能作为 AMP 包（torch.cuda.amp）移入了 PyTorch 核心。`torch.cuda.amp` 比 `apex.amp` 更灵活、更直观。`torch.cuda.amp` 能够解决 `apex.amp` 的一些已知痛点：

保证 PyTorch 版本兼容性，因为它本身就是 PyTorch 的一部分
无需构建扩展
Windows 支持
检查点的位精确保存/恢复
DataParallel 和进程内模型并行（尽管我们仍建议使用每进程一个 GPU 的 torch.nn.DistributedDataParallel 作为性能最佳的方法）
梯度惩罚（二次反向传播）
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()

性能基准

在本节中，我们将讨论在使用 NGC 中的 NVIDIA pytorch:20.06-py3 容器运行深度学习工作负载时，AMP 在最新 NVIDIA GPU A100 以及上一代 V100 GPU 上的混合精度训练的准确性和性能。混合精度性能将与 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)

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

训练精度：NVIDIA DGX-1 (8x V100 16GB)

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

加速性能：

NVIDIA V100 上的 FP16 vs. V100 上的 FP32

对于 V100 上的 DL 训练，AMP 与 FP16 结合是性能最佳的选择。在表 1 中，我们可以观察到对于各种模型，V100 上的 AMP 相对于 V100 上的 FP32 提供了 1.5 倍到 5.5 倍的加速，同时收敛到相同的最终精度。

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

NVIDIA A100 上的 FP16 与 V100 上的 FP16

对于 A100 上的 DL 训练，结合 FP16 的 AMP 仍然是性能最佳的选择。在图 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 的模型和模型脚本，并开始进行实验。原生 AMP 支持的 NVIDIA PyTorch 可从 PyTorch NGC 容器 20.06 版本获得。我们强烈建议现有的 `apex.amp` 客户过渡到使用最新 PyTorch 1.6 版本中提供的 PyTorch Core 中的 `torch.cuda.amp`。

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