tf.keras.mixed_precision.policy必须在模型构建前设置，因混合精度依赖计算图构建时确定dtype；延迟设置将导致权重仍为float32、梯度缩放失效，甚至报错。
为什么 tf.keras.mixed_precision.Policy 必须在模型构建前设置
混合精度训练依赖于计算图构建阶段就确定每层的 dtype 行为，如果在模型编译或训练之后再启用策略，tf.keras 不会自动重置已有层的权重 dtype 或重新绑定梯度缩放逻辑，导致实际仍以 float32 运行，甚至引发 InvalidArgumentError: Cannot assign a device for operation 类错误。
正确做法是：在导入 TensorFlow 后、定义模型前，立即配置全局策略：
import tensorflow as tf
policy = tf.keras.mixed_precision.Policy('mixed_float16')
tf.keras.mixed_precision.set_global_policy(policy)
- 必须在
tf.keras.Model子类定义或tf.keras.Sequential构建之前调用set_global_policy - 若使用自定义训练循环，还需手动添加
tf.keras.mixed_precision.LossScaleOptimizer包装器 - GPU 上默认启用
mixed_float16；TPU 则应使用mixed_bfloat16，否则会报NotImplementedError: bfloat16 is not supported on this device
tf.keras.mixed_precision.LossScaleOptimizer 在自定义训练循环中不可省略
启用 mixed_float16 后，梯度容易因数值下溢（underflow）变成 0，尤其在小 loss 或深层网络中。Loss scaling 是必须步骤，而 LossScaleOptimizer 就是这个机制的封装 —— 它不是可选项，是保障收敛性的刚需组件。
典型错误写法是只设 policy、不包装 optimizer：
# ❌ 错误：没有 loss scaling，训练几步后 loss 突然变为 nan optimizer = tf.keras.optimizers.Adam() model.compile(optimizer=optimizer, loss='sparse_categorical_crossentropy')
正确写法（自定义循环）：
optimizer = tf.keras.optimizers.Adam()
optimizer = tf.keras.mixed_precision.LossScaleOptimizer(optimizer)
# ……然后在 @tf.function 训练步中：
with tf.GradientTape() as tape:
predictions = model(x, training=True)
loss = loss_fn(y, predictions)
scaled_gradients = tape.gradient(loss, model.trainable_variables)
optimizer.apply_gradients(zip(scaled_gradients, model.trainable_variables))
-
LossScaleOptimizer会自动在前向时放大 loss，在反向时缩放梯度，无需手动干预 - 若用
model.fit()，Keras 内部已集成该逻辑，此时可跳过显式包装，但前提是 policy 已提前设置 - 检查是否生效：打印
optimizer.loss_scale，正常值应为DynamicLossScale实例，而非FixedLossScale
验证混合精度是否真正生效的三个关键检查点
光看代码设了 policy 不代表 GPU 真的跑在 float16 —— 常见失效原因包括：模型中有强制 float32 的操作、数据未适配、或部分层被忽略 dtype 推导。
快速验证方法：
- 检查层输出 dtype：
model.layers[0](tf.random.normal((1,28,28,1))).dtype应返回tf.float16（除最后几层如 Dense 分类头通常保持 float32） - 监控 GPU 显存：启用 mixed precision 后，同等 batch size 下显存占用应下降约 30%~40%，否则大概率未生效
- 查看
tf.config.list_physical_devices('GPU')输出，确认 CUDA/cuDNN 版本支持 FP16（TensorFlow 2.9+ 要求 CUDA 11.2+、cudnn 8.1+）
特别注意：tf.nn.softmax 和 tf.keras.activations.softmax 在 float16 下不稳定，建议显式 cast 回 float32 再 softmax，或改用 tf.keras.layers.Softmax(dtype='float32')。
BatchNorm 和 Dropout 层在 mixed_float16 下的隐含风险
这两个层内部状态（running mean/variance、mask）默认继承输入 dtype，但在 mixed_float16 模式下，它们若以 float16 更新，会导致统计量漂移、训练不稳定，甚至 divergence。
TensorFlow 2.x 默认已对 tf.keras.layers.BatchNormalization 和 Dropout 做了 dtype 保护 —— 它们会自动将内部变量保持为 float32，仅输入/输出走 float16。但前提是：你没手动指定 dtype='float16'。
- ✅ 正确（依赖默认行为）：
tf.keras.layers.BatchNormalization() - ❌ 危险（破坏保护）：
tf.keras.layers.BatchNormalization(dtype='float16')—— 会强制所有内部计算走 float16，大概率崩溃 - 自定义层中若用到
tf.nn.batch_normalization等底层 API，必须显式传入mean和variance的 float32 张量
最稳妥的做法是：不碰这些层的 dtype 参数，让 Keras 自动处理；如有疑问，直接打印 layer.moving_mean.dtype 验证。
Python免费学习笔记（深入）：立即使用
在学习笔记中，你将探索 Python 的核心概念和高级技巧！
