混合精度

什么是混合精度训练？

混合精度训练是在模型训练期间使用较低精度操作（float16 和 bfloat16），以使其运行更快并使用更少内存。使用混合精度可以在现代 GPU 上将性能提升 3 倍以上，在 TPU 上提升 60%。

目前，大多数模型使用 float32 数据类型，它占用 32 位内存。然而，还有两种较低精度的数据类型：float16 和 bfloat16，它们都只占用 16 位内存。像 Google TPU 和 NVIDIA GPU 这样的现代加速器可以更快地运行 16 位数据类型操作，因为它们拥有专门的硬件来执行 16 位计算，而且 16 位数据类型可以更快地从内存读取。因此，在这些设备上应尽可能使用这些较低精度的数据类型。

然而，变量存储（以及某些敏感计算）仍应使用 float32 以保持数值稳定性。通过尽可能使用 16 位精度并保持模型的某些关键部分使用 float32，模型在训练和使用 32 位精度时都将运行得更快。

在 Keras 中使用混合精度训练

Keras 层或模型使用的精度策略由 keras.mixed_precision.DTypePolicy 实例控制。每个层都有自己的 DTypePolicy。您可以通过 dtype 参数在单个层上设置它（例如 MyLayer(..., dtype="mixed_float16")），或者通过工具函数 keras.mixed_precision.set_global_policy 设置一个全局值，供所有层默认使用。

通常，要在 GPU 上开始使用混合精度，您只需在程序开始时调用 keras.mixed_precision.set_global_policy("mixed_float16")。在 TPU 上，您应调用 keras.mixed_precision.set_global_policy("mixed_bfloat16")。

API 文档

混合精度策略 API

支持的硬件

虽然混合精度可以在大多数硬件上运行，但它只会在较新的 NVIDIA GPU 和 Google TPU 上加速模型。NVIDIA GPU 支持混合使用 float16 和 float32，而 TPU 支持混合使用 bfloat16 和 float32。

在 NVIDIA GPU 中，计算能力为 7.0 或更高的 GPU 将从混合精度中获得最大的性能提升，因为它们具有称为 Tensor Cores 的特殊硬件单元，可以加速 float16 矩阵乘法和卷积。较旧的 GPU 使用混合精度无法获得数学性能上的好处，但内存和带宽的节省可以实现一定的加速。您可以在 NVIDIA 的 CUDA GPU 网页上查找您 GPU 的计算能力。将从混合精度中获益最多的 GPU 示例包括 RTX GPU、V100 和 A100。

即使在 CPU 和较旧的 GPU 上，尽管预期不会有加速，混合精度 API 仍可用于单元测试、调试或仅仅试用 API。然而，在 CPU 上，混合精度将运行得明显更慢。