Keras mixed_float16 训练出现 NaN 怎么区分 loss scale 和输入范围

Keras 使用 mixed_float16 后训练出现 NaN，不能先入为主地归因于 loss scale。更稳妥的判断顺序是：先用同一份数据跑一个 float32 基线，再检查输入是否有限、范围是否过大，最后比较 mixed precision 下的前向损失和梯度是否先变成非有限值。

官方地址：https://keras.io/api/mixed_precision/

这个顺序很重要：如果输入在进入模型前就包含 NaN 或 Inf，调 loss scale 没有意义；如果只有 mixed_float16 对照失败，才需要继续看敏感计算、损失缩放、学习率和梯度。Keras 官方文档把低精度计算与 float32 变量、敏感计算分开处理，本文沿着这个边界给出排查路径。

先把 mixed_float16 的变化边界说清楚

Keras 的 dtype policy 决定层在计算时使用什么精度，以及变量保存采用什么精度。设置 mixed_float16 后，常见的结构是：部分算子用 float16，变量仍保持 float32，而最终损失、归一化、指数或归约等敏感位置需要特别留意。

因此 NaN 的来源至少有三个边界：

• 数据边界：输入已经含有 NaN、Inf，或者数值范围远大于模型预期。
• 计算边界：float16 的表示范围、指数运算、平方和或归约让前向结果溢出。
• 优化边界：损失经过梯度计算后，梯度或更新量变成非有限值，损失缩放没有消除这个问题。

图中的连线表示概念之间的依赖关系，不表示某次训练已经执行或某个值一定会溢出。排查时要把这三个边界分别记录，避免看到第一个 NaN 就修改一整套配置。

先做 float32 基线和输入有限性检查

最有价值的第一步是保留数据、模型结构、批大小和优化器，暂时把策略切回 float32。如果基线同样出现 NaN，问题优先落在数据、标签、损失公式或学习率；如果基线正常而 mixed_float16 失败，才把注意力放到低精度计算路径和损失缩放。

import numpy as np
import keras
from keras import layers

# 先用 float32 建立对照，保持数据和模型结构不变。
keras.mixed_precision.set_global_policy("float32")

inputs = keras.Input(shape=(128,), dtype="float32", name="features")
# 输入层只负责接收已经完成基本清洗的特征。
x = layers.Dense(256, activation="relu", name="hidden")(inputs)
# 输出保持 float32，避免损失读取低精度预测值时扩大数值误差。
outputs = layers.Dense(1, dtype="float32", name="prediction")(x)
model = keras.Model(inputs, outputs)
model.compile(
    # 先使用较保守的学习率，便于把数据问题和更新过大区分开。
    optimizer=keras.optimizers.Adam(learning_rate=1e-3),
    loss=keras.losses.MeanSquaredError(),
)

def report_array(name, value):
    # 只观察有限性和范围，不把单个极值误认为 loss scale 证据。
    values = np.asarray(value)
    finite = np.isfinite(values)
    print(name, "dtype=", values.dtype, "finite=", bool(finite.all()))
    if finite.any():
        print("range=", float(values[finite].min()), float(values[finite].max()))

# 在送入 fit 前检查特征和标签，NaN/Inf 不应进入后续判断。
report_array("features", x_train)
report_array("labels", y_train)

如果这里的 finite 已经是 False，先处理数据清洗、除零、对数输入和标签生成逻辑。若输入全部有限但范围跨越很多数量级，也要先确认标准化、裁剪或单位换算，而不是立即调 loss scale。

用输入范围与前向损失定位溢出点

float16 并不等于所有中间值都会自动安全转换。尤其是平方损失、指数、softmax 前的 logits、长序列归约和自定义层中的除法，都可能把一个看似有限的输入放大到不可表示的范围。

可以用同一批样本做三个对照：

1. 原始输入送入 float32 基线，记录预测和损失是否有限。
2. 对输入做与业务含义一致的归一化，再观察损失是否恢复有限。
3. 保持归一化不变，只切换 mixed_float16，判断差异是否只在低精度路径出现。

这里的关键不是追求某一个“安全最大值”，而是把数据单位和模型期望对应起来。例如像素值是否仍是 0–255、金额是否混用了分和元、标准差是否接近零，都会比 loss scale 更早造成异常。

# 这个片段用于在进入模型前记录输入范围，避免把单位错误误判成精度问题。
features = np.asarray(x_train, dtype="float32")
labels = np.asarray(y_train, dtype="float32")

# 先确认数组本身有限，再观察分位数和绝对最大值。
if not np.isfinite(features).all() or not np.isfinite(labels).all():
    raise ValueError("输入或标签含有 NaN/Inf，请先修复数据管道")

print("feature p01/p99:", np.percentile(features, [1, 99]))
print("feature abs max:", np.max(np.abs(features)))
print("label abs max:", np.max(np.abs(labels)))

# 归一化应与训练和推理保持同一套统计口径。
mean = features.mean(axis=0, keepdims=True)
std = features.std(axis=0, keepdims=True)
safe_std = np.maximum(std, 1e-6)  # 防止常量列除以零。
normalized = (features - mean) / safe_std
print("normalized finite:", bool(np.isfinite(normalized).all()))

如果原始输入的范围异常，修复数据管道后再复现一次；如果输入范围合理、float32 正常而 mixed_float16 的前向损失先变为 NaN，可以继续检查损失中的敏感运算和输出层 dtype。最终输出保持 float32 是常见的稳定化边界，但不要把所有层都强制回退，否则失去混合精度的目标。

再判断 loss scale 与梯度是否相关

损失缩放的作用是让较小的梯度在低精度计算阶段不容易下溢，再在更新前还原尺度。它解决的是梯度表示和更新稳定性问题，不能修复已经进入模型的 NaN，也不能替代输入归一化。

可以用下面的证据来区分：

如果框架已经为标准训练路径提供损失缩放，优先使用框架的默认机制，并在自定义 train_step、自定义 loss 或手动梯度流程中确认缩放和反缩放没有被绕过。不要同时修改 dtype policy、学习率、损失公式和输入预处理，否则无法知道是哪一项真正起作用。

按证据选择迁移修复

确定原因后，修复应尽量小：

• 数据范围问题：在数据管道中完成稳定、可复现的归一化；训练和推理使用同一套统计量。
• 敏感计算问题：只把输出层、损失相关计算或确实容易溢出的自定义层设为 dtype="float32"。
• 梯度问题：确认优化器的损失缩放路径没有被自定义训练循环绕过，再考虑降低学习率或裁剪梯度。
• 基线也不稳定：先修复损失函数中的除零、log(0)、空归约和标签异常，不要把回退 float32 当成根治。

一个实用的迁移顺序是：先用 float32 建立可收敛基线，再恢复 mixed_float16；保持输出和损失的 float32 边界；最后才微调学习率、批大小或梯度裁剪。每次只改变一个变量，并保留同一批样本的有限性记录。

用最小验证固定排查结论

最后保留一份小而稳定的诊断记录即可，不需要把训练跑成大型实验。记录同一批样本在两种 policy 下的输入范围、前向损失、梯度有限性和第一处非有限值位置。

# 这个检查器用于比较两种 policy 的关键数值证据。
def finite_summary(name, value):
    values = np.asarray(value)
    finite = np.isfinite(values)
    return {
        "name": name,
        "dtype": str(values.dtype),
        "finite": bool(finite.all()),
        # 没有有限值时不计算 min/max，避免诊断器本身再次产生异常。
        "min": float(values[finite].min()) if finite.any() else None,
        "max": float(values[finite].max()) if finite.any() else None,
    }

def compare_policy(policy_name, model, batch_x, batch_y):
    # 只做前向和损失观察；自定义训练循环需另行记录梯度。
    prediction = model(batch_x, training=False)
    loss_value = keras.losses.mean_squared_error(batch_y, prediction)
    return {
        "policy": policy_name,
        "prediction": finite_summary("prediction", prediction),
        "loss": finite_summary("loss", loss_value),
    }

# 先分别建立 float32 和 mixed_float16 模型，再使用完全相同的 batch。
print(compare_policy("float32", float_model, batch_x, batch_y))
print(compare_policy("mixed_float16", mixed_model, batch_x, batch_y))

这段代码是排查框架，不是某台机器上的运行结果。若两种 policy 的前向都有限、只有反向更新后出现 NaN，就把观察点放到梯度和优化器；若 mixed_float16 的前向已经异常，就回到输入范围、敏感层和损失计算。

几个容易混淆的问题

只要出现 NaN 就应该关闭 mixed_float16 吗？

不应该。先看 float32 基线和输入有限性。如果基线也失败，关闭 mixed precision 只是把问题隐藏在更宽的数值范围里；如果只有低精度路径失败，再针对敏感计算和缩放边界修复。

输入已经是 float32，为什么还会受 mixed_float16 影响？

输入 dtype 只是入口信息，层的 policy 仍可能让中间计算使用低精度。真正需要观察的是归一化、激活、损失和输出的边界，而不是只打印输入 dtype。

降低学习率能证明是 loss scale 问题吗？

不能。降低学习率可能让更新变小，但它同时改变了优化动力学。只有在输入和前向损失稳定、梯度或更新阶段出现非有限值时，才把学习率与损失缩放放进同一组证据里比较。

总结

Keras mixed_float16 出现 NaN 时，先做 float32 对照和输入有限性检查，再判断 NaN 出现在前向损失、梯度还是参数更新阶段。输入范围异常属于数据边界，敏感算子溢出属于计算边界，梯度非有限才值得深入看 loss scale。按这个顺序每次只改一个因素，才能得到可复用的迁移结论。