大多数深度学习框架，包括 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 和进程内模型并行（尽管我们仍然推荐torch.nn.DistributedDataParallel，每个进程一个 GPU，作为性能最佳的方法）
梯度惩罚（二次反向传播）
torch.cuda.amp.autocast() 在未启用区域无效，因此它应该能够轻松处理以前在多次调用 apex.amp.initialize() 时（包括交叉验证）遇到的问题。在同一个脚本中进行多次收敛运行时，每次都应该使用一个新的 GradScaler 实例，但 GradScaler 轻巧且自包含，因此这不是问题。
稀疏梯度支持

随着 AMP 被添加到 PyTorch 核心，我们已经开始弃用 apex.amp。我们已将 apex.amp 移至维护模式，并将继续支持使用 apex.amp 的客户。但是，我们强烈建议 apex.amp 客户过渡到使用 PyTorch Core 中的 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 A100 GPU 和上一代 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

加速性能：

NVIDIA V100 上的 FP16 vs. V100 上的 FP32

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

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

NVIDIA A100 上的 FP16 vs. V100 上的 FP16

AMP 与 FP16 仍然是 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 的模型和模型脚本。支持原生 AMP 的 NVIDIA PyTorch 可从 PyTorch NGC 容器版本 20.06 获得。我们强烈鼓励现有 apex.amp 客户过渡到使用最新 PyTorch 1.6 版本中提供的 PyTorch Core 中的 torch.cuda.amp。

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