TensorFlow 的自动混合精度，可提升在 NVIDIA GPU 上 AI 训练速度

文 / NVIDIA

混合精度训练利用半精度来加速训练，在某些情况下，能够获得与使用相同超参数进行单精度训练相同的精度。内存要求也因此降低，从而可以运行更大的模型和进行小批量处理。

启用混合精度需要完成两个步骤：移植模型，以在适当的时候使用半精度数据类型；使用损失缩放技术保留小的梯度值。我们为 TensorFlow 引入自动混合精度功能（现可在 1.x 中使用，不久后即可在 2.x 中使用），该功能使用 NVIDIA Volta 和 Turing GPU 中的 Tensor Core 进行修改，以提升训练性能。自动混合精度在 TensorFlow 内部通过 NVIDIA NGC Container 中的一个环境变量完成这两个步骤，并在必要时采取更精细的控制。

为现有的 TensorFlow 模型脚本启用此功能需要设置一个环境变量，或者只需更改几行代码。我们发现对于数学运算量大的模型，运算速度提高了 3 倍，但具体能获得多少程度的加速取决于模型架构。现在，TensorFlow 容器内提供自动混合精度功能，而 TensorFlow 容器可在 NVIDIA NGC 容器注册表 中取得（https://ngc.nvidia.com/catalog/containers/nvidia:tensorflow）。

要在容器中启用此功能，只需设置一个环境变量即可：

export TF_ENABLE_AUTO_MIXED_PRECISION=1

或者，您也可以在 TensorFlow Python 脚本中设置该环境变量：

os.environ[‘TF_ENABLE_AUTO_MIXED_PRECISION’] = ‘1’

启用混合精度后，您可通过以下方式进一步实现加速：

• 启用 TensorFlow XLA 编译器，但请注意，Google 仍将 XLA 列为实验性工具

• 增加小批量大小。较大的小批量通常可以提高 GPU 的利用率，而混合精度最多可以将小批量大小增加到原来的 2 倍

加速是采用单精度和自动混合精度执行固定次数训练的时间比。

每个模型的训练次数与文献或常见做法相符（并且可以确认的是，两种训练的模型精度相同）。

批量大小测量如下：rn50 (v1.5)：FP32 为 128，AMP+XLA 为 256；ssd-rn50-fpn-640：FP32 为 8，AMP+XLA 为 16；ncf：FP32 和 AMP+XLA 为 1M；bert-squadqa：FP32 为 4，AMP+XLA 为 10；gnmt：FP32 为 128，AMP 为 192。

您可以在 https://github.com/NVIDIA/DeepLearningExamples/tree/master/TensorFlow上找到所有模型，但 ssd-rn50-fpn-640 除外，请前往以下位置获取该模型： https://github.com/tensorflow/models/tree/master/research/object_detection。除了 bert-squadqa 的性能数据是在 1xV100 (32GB) 上采集之外，所有性能数据均在 1xV100 (16GB) 上采集。

可用性

NVIDIA 经优化的 TensorFlow 19.03 NVIDIA NGC Container 中提供自动混合精度功能。我们还与 Google 的 TensorFlow 团队紧密合作，致力于将此功能直接并入 TensorFlow 框架核心。

您还可以在 NVIDIA NGC 模型脚本注册表 或 GitHub 中找到我们用来生成上述性能图的示例训练文本。

注：NVIDIA NGC 模型脚本注册表 链接

https://ngc.nvidia.com/catalog/model-scripts

GitHub 链接

https://github.com/NVIDIA/DeepLearningExamples/tree/master/TensorFlow

请试用 NVIDIA 经优化的 TensorFlow 容器，开始体验自动混合精度。您可以随时在 TensorFlow 论坛中向我们的团队提供反馈或提出问题。

其他资源

• NVIDIA GPU 上的自动混合精度

  （https://developer.nvidia.com/automatic-mixed-precision）

• 开发者博客：TensorFlow 中适用于 NVIDIA Tensor Core 架构的自动混合精度

  （https://devblogs.nvidia.com/nvidia-automatic-mixed-precision-tensorflow/）

• NVIDIA GPU 上的 Tensor Core

  （https://developer.nvidia.com/tensor-cores）

更多 AI 相关阅读：