TensorFlow 2.x混合精度训练需先设全局策略mixed_float16，再用LossScaleOptimizer包装优化器；输入保持float32，层自动转float16计算但权重保留float32；自定义层/loss须适配dtype，避免NaN和显存异常。
TensorFlow 2.x 中启用混合精度训练的正确方式
TensorFlow 2.3+ 原生支持混合精度（AMP），但不是靠手动 cast，而是通过 tf.keras.mixed_precision.Policy 和全局策略控制。直接在模型层加 tf.cast 或用 tf.float16 初始化权重，反而会破坏自动梯度缩放逻辑，导致 NaN 梯度或训练崩溃。
核心操作只有两步：设置全局策略 + 启用损失缩放器。其余由 Keras 自动处理。
- 必须在构建模型前调用
tf.keras.mixed_precision.set_global_policy("mixed_float16") - 必须使用
tf.keras.mixed_precision.LossScaleOptimizer包装优化器（TF 2.9+ 已默认集成进model.compile，但显式包装更可控） - 输入数据仍应为
float32—— 策略会自动在计算密集层（如 Dense、Conv2D）内部转为float16，但保留float32的主权重和累加器
为什么 model.compile(..., mixed_precision=True) 不生效？
这个参数根本不存在。常见误操作是查文档时混淆了 PyTorch 的 amp.autocast 接口，或误读旧版 TF 文档中未发布的草案。TensorFlow 没有 mixed_precision=True 这类 compile 参数。
真正起作用的是策略 + 优化器组合。如果你发现 loss 不下降或 nan 突然爆发，大概率是漏掉了 LossScaleOptimizer，或策略设置太晚（比如在 model = ... 之后才 set_global_policy）。
立即学习“Python免费学习笔记（深入）”；
快速生成专业的 Python 脚本和应用代码。一键创建完整项目结构，支持CLI、API、爬虫、Bot、Django等多种项目类型，包含完整的项目结构、配置文件、依赖管理、测试、README和文档。
- 策略必须在任何 Keras 层实例化前设置，否则已有层不会继承新策略
- 验证是否生效：打印某一层（如
model.layers[0]）的compute_dtype和variable_dtype，应分别为float16和float32 - 用
tf.debugging.enable_check_numerics()可快速定位哪步产生 inf/nan，常出问题的是自定义 loss 或非标准激活函数（如未适配float16范围的tf.math.log）
自定义层/loss 中如何安全使用混合精度？
自定义 tf.keras.layers.Layer 必须显式声明 self._compute_dtype 和 self._variable_dtype，否则会退回到全 float32。而自定义 loss 函数若含数值敏感操作（如 tf.clip_by_value、tf.math.softplus），需主动升回 float32 再计算。
class SafeCustomLayer(tf.keras.layers.Layer):
def __init__(self, **kwargs):
super().__init__(**kwargs)
# 必须手动同步策略
self._compute_dtype = tf.keras.mixed_precision.global_policy().compute_dtype
self._variable_dtype = tf.keras.mixed_precision.global_policy().variable_dtype
<p>def custom_loss(y_true, y_pred):</p><h1>y_pred 是 float16，直接 log 容易 underflow</h1><pre class='brush:python;toolbar:false;'>y_pred_f32 = tf.cast(y_pred, tf.float32)
return tf.keras.losses.categorical_crossentropy(y_true, y_pred_f32)- 不要在
call()中硬编码tf.float16；用self._compute_dtype动态获取 - 所有涉及指数、对数、除法、平方根的操作，优先在
float32下完成，再 cast 回输出 dtype - 避免在
float16下做tf.reduce_sum大量小数值 —— 累加误差会被放大，建议中间用tf.float32累加
GPU 显存没降反升？检查这三处
混合精度本应降低显存占用约 30%，但如果显存不降甚至升高，通常是策略未生效或额外开销掩盖了收益。
- 确认 GPU 是否支持 FP16 加速：运行
nvidia-smi --query-gpu=name --format=csv，Tesla V100 / A100 / RTX 30xx+ 才有原生 Tensor Core 支持；GTX 10xx 系列仅软件模拟，无加速且易溢出 - 检查是否启用了
tf.data.Dataset.cache()或prefetch(tf.data.AUTOTUNE)—— 这些预处理会在 CPU 内存中缓存原始float32数据，显存看似没减，实则 CPU 内存涨了 - 查看
tf.config.optimizer.set_jit(True)是否开启：XLA 编译可进一步提升混合精度效率，尤其对小 batch 场景，但首次运行会多占显存编译
混合精度不是开关一按就完事，它把数值稳定性压力转移给了 loss scale 动态调整和 layer 实现细节。最容易被忽略的是：自定义 ops 和第三方库（如 tensorflow-addons）往往不兼容 mixed_float16 策略，需要单独验证 dtype 行为。
