汤不热吧引言
在训练大型深度学习模型时，显存（VRAM）往往是最大的瓶颈之一。TensorFlow 2.x 引入了强大的混合精度训练（Mixed Precision Training）功能，允许我们在不牺牲模型精度的情况下，大幅减少显存占用并提高训练速度。
混合精度训练的核心思想是利用 FP16（半精度浮点数，16位）来存储模型的激活值和梯度，而将模型的权重（Weights）和优化器状态保留在 FP32（单精度浮点数，32位）。由于 FP16 只占用 FP32 一半的存储空间，理论上可以直接节省约 50% 的显存。
本文将聚焦于如何使用 Keras Policy 接口来启用这一功能，并提供一个完整的实操示例。
技术核心：Keras Mixed Precision Policy
TensorFlow Keras 使用 tf.keras.mixed_precision.Policy 来管理计算精度。对于现代 GPU（如 Volta、Turing 或 Ampere 架构），推荐使用 ‘mixed_float16’ 策略。
Policy **’mixed_float16′ 工作原理：**
- 变量/权重 (FP32)： 模型权重和变量默认仍保持 FP32，以保证数值稳定性。
- 计算/激活 (FP16)： 模型层中的计算（如矩阵乘法）和中间激活值使用 FP16。
- 损失缩放 (Loss Scaling)： 启用 FP16 后，由于其数值范围较小，可能会导致小梯度在反向传播时下溢（underflow）为零。因此，mixed_float16 策略会自动要求使用 tf.keras.optimizers.LossScaleOptimizer 来解决这一问题。
实操步骤与代码示例
我们将使用一个简单的 CNN 模型在 MNIST 数据集上进行演示。
1. 环境准备
确保您的 TensorFlow 版本大于 2.4，并且拥有支持 Tensor Cores 的 NVIDIA GPU。
import tensorflow as tf
from tensorflow.keras.layers import Conv2D, Dense, Flatten
from tensorflow.keras.models import Sequential
from tensorflow.keras.optimizers import Adam
# 检查GPU是否可用
if not tf.config.list_physical_devices('GPU'):
print("WARNING: No GPU detected. Mixed precision requires a compatible GPU.")
# 检查TensorFlow版本
print(f"TensorFlow version: {tf.__version__}")
2. 设置混合精度 Policy
这是启用混合精度的关键步骤。我们设置全局策略为 ‘mixed_float16’。
# 核心步骤：设置混合精度策略
from tensorflow.keras.mixed_precision import Policy
policy = Policy('mixed_float16')
tf.keras.mixed_precision.set_global_policy(policy)
print(f"Global policy set to: {policy.name}")
print(f"Compute dtype: {policy.compute_dtype}") # 预期输出 float16
print(f"Variable dtype: {policy.variable_dtype}") # 预期输出 float32
3. 构建模型
当全局 Policy 设置完成后，模型中所有支持混合精度的层（如 Dense, Conv2D）将自动使用 FP16 作为计算数据类型。
def create_cnn_model():
model = Sequential([
Conv2D(32, 3, activation='relu', input_shape=(28, 28, 1)),
Flatten(),
Dense(10)
])
return model
model = create_cnn_model()
4. 数据加载与训练
由于我们使用了 ‘mixed_float16’ 策略，优化器必须使用 LossScaleOptimizer 进行封装，以防止梯度下溢。
# 准备数据
(x_train, y_train), _ = tf.keras.datasets.mnist.load_data()
x_train = x_train[..., tf.newaxis].astype('float32') / 255.0
# 优化器必须使用 LossScaleOptimizer 封装
opt = Adam(learning_rate=1e-3)
opt = tf.keras.mixed_precision.LossScaleOptimizer(opt)
model.compile(
optimizer=opt,
loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
metrics=['accuracy']
)
print("\n--- 开始混合精度 (FP16) 训练 ---")
# 训练一小段时间
model.fit(x_train, y_train, batch_size=256, epochs=2, verbose=2)
print("\n混合精度训练完成，显存占用显著降低。")
为什么混合精度能节省显存？
混合精度节省显存的原理非常直观：
- 激活值存储减半： 在前向传播过程中，模型的激活值（Activation）是主要的显存消耗者之一。当计算类型从 FP32 切换到 FP16 时，每个激活值占用的空间从 4 字节降至 2 字节。
- 梯度存储减半： 在反向传播过程中，梯度也需要存储在显存中。同样，如果梯度使用 FP16 存储，所需的空间也减半。
对于批量大小（Batch Size）较大或层数较深的场景，这种 50% 的内存节省允许用户使用更大的 Batch Size，从而加速收敛并更好地利用 GPU 资源。