TensorFlow 混合精度下自定义损失函数为何出现 dtype 不匹配？

2026-11-04数据小陈

在 TensorFlow 的 mixed_float16 策略下，自定义损失函数接收的 y_true 和 y_pred 可能分别是 float32 和 float16。示例中的 y_true - y_pred 会因两者类型不同而触发 Sub 操作错误。

为什么会出现类型不匹配

混合精度策略通常让模型层使用 float16 计算，同时让变量保持为 float32。训练标签通常仍为 float32，模型输出则可能是 float16。因此，损失函数需要处理输入类型不同的情况。

在损失函数中统一使用 float32

对数值敏感的损失计算，可以先将预测值和标签都转换为 float32，再计算差值和均值：

import tensorflow as tf
from tensorflow.keras import mixed_precision

mixed_precision.set_global_policy("mixed_float16")

def custom_mse(y_true, y_pred):
    y_true = tf.cast(y_true, tf.float32)
    y_pred = tf.cast(y_pred, tf.float32)
    diff = y_true - y_pred
    return tf.reduce_mean(tf.square(diff))

这样，减法两侧的 dtype 就一致了，损失计算也会在 float32 中进行。将这个损失函数用于原来的 model.compile(optimizer=tf.keras.optimizers.Adam(), loss=custom_mse)，可以避免示例中的减法类型冲突。

可选做法：让输出层返回 float32

如果需要模型最后一层输出 float32，可以显式指定该层的 dtype：

model = tf.keras.Sequential([
    tf.keras.layers.Dense(64, activation="relu"),
    tf.keras.layers.Dense(1, dtype="float32"),
])

这样最后一层会输出 float32。自定义损失函数仍可显式转换输入类型，以免依赖输出层的配置。

注意事项

原问题片段中的 diff = y_true - y_pred 和错误信息都指向 Sub 运算。应检查损失函数中参与运算的张量类型，并将它们统一。具体 dtype 行为可能因 TensorFlow 和 Keras 版本以及模型配置而异。遇到类似问题时，可以查看损失函数中的 y_true.dtype 和 y_pred.dtype，再选择统一的计算类型。