TensorFlow2.15训练的分类模型无法在2.16导入，报reduction参数错误

问题根源 #

TensorFlow 2.16移除了损失函数reduction参数的'auto'取值，而2.15版本训练的模型保存时可能将损失函数的reduction设为'auto'，导致导入时触发参数校验错误。

解决方案 #

方法1：在TF2.15环境中修复模型后重新导出

如果能访问TF2.15环境，直接修改模型的损失函数参数后重新保存：

import tensorflow as tf

# 加载原模型
model = tf.keras.models.load_model("your_trained_model_path")

# 显式指定损失函数的reduction为合法值（'sum_over_batch_size'等价于原'auto'的默认行为）
loss_fn = tf.keras.losses.CategoricalCrossentropy(reduction='sum_over_batch_size')
# 若模型使用的是其他损失函数（如SparseCategoricalCrossentropy），替换为对应类即可

# 重新编译并保存模型
model.compile(optimizer=model.optimizer, loss=loss_fn, metrics=model.metrics)
model.save("fixed_model_path")

将修复后的模型在TF2.16中导入即可正常使用。

方法2：在TF2.16环境中兼容处理

无法回到TF2.15环境时，可通过包装损失函数类，自动将'auto'映射为合法参数：

import tensorflow as tf

def wrap_loss(loss_class):
    class WrappedLoss(loss_class):
        def __init__(self, *args, **kwargs):
            # 将reduction='auto'替换为等价的'sum_over_batch_size'
            if kwargs.get("reduction") == "auto":
                kwargs["reduction"] = "sum_over_batch_size"
            super().__init__(*args, **kwargs)
    return WrappedLoss

# 根据模型使用的损失函数替换对应类，示例为CategoricalCrossentropy
tf.keras.losses.CategoricalCrossentropy = wrap_loss(tf.keras.losses.CategoricalCrossentropy)

# 现在导入模型
model = tf.keras.models.load_model("your_trained_model_path")

补充说明 #

'sum_over_batch_size'是TensorFlow中损失函数reduction='auto'的默认等价行为，适合绝大多数分类任务场景。

内容的提问来源于stack exchange，提问作者Thriver