Keras 3 API 文档 / 混合精度

混合精度

什么是混合精度训练?

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

目前,大多数模型使用 float32 数据类型,它占用 32 位内存。然而,还有两种较低精度的数据类型:float16bfloat16,它们都只占用 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 上,混合精度将运行得明显更慢。

您可以使用以下命令检查您的 GPU 类型

nvidia-smi -L