► Keras 3 API文档 / 混合精度

混合精度

什么是混合精度训练？

混合精度训练是指在训练模型时使用较低精度的运算（float16 和 bfloat16）来提高运行速度和减少内存占用。使用混合精度可以在现代 GPU 上将性能提高 3 倍以上，在 TPU 上提高 60%。

如今，大多数模型都使用 float32 数据类型，它占用 32 位内存。然而，还有两种较低精度的 16 位数据类型，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 上，混合精度运行速度会明显变慢。