怎样在Python中通过TensorFlow实现自动混合精度AMP-Python教程-PHP中文网

tensorflow 2.x混合精度训练需先设全局策略mixed_float16，再用lossscaleoptimizer包装优化器；输入保持float32，层自动转float16计算但权重保留float32；自定义层/loss须适配dtype，避免nan和显存异常。

怎样在python中通过tensorflow实现自动混合精度amp

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-script-generator

快速生成专业的 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 行为。