TensorFlow2.3.0加载.h5模型报错：__init__()存在意外关键字参数axis

错误原因 #

报错核心是训练模型时使用的TensorFlow版本与当前加载环境的2.3.0版本不兼容：高版本TensorFlow的内置损失函数（如CategoricalCrossentropy、SparseCategoricalCrossentropy）初始化支持axis参数，保存模型时该参数会写入模型的训练配置中，而2.3.0版本的损失函数没有这个初始化参数，反序列化损失配置时就会抛出参数不匹配的错误。

开头的CPU指令集优化提示是正常的性能提示，不影响模型运行，无需特殊处理。

解决方案 #

• 方案1：仅做推理不需要后续训练时，加载模型关闭编译配置即可跳过损失函数反序列化，修改加载代码为：

# 新增compile=False参数
age_model = keras.models.load_model(age_model_path, compile=False)
gender_model = keras.models.load_model(gender_model_path, compile=False)

• 方案2：需要保留模型编译配置用于后续训练时，传入自定义的损失反序列化规则，过滤掉axis参数：

from tensorflow.keras.losses import SparseCategoricalCrossentropy
# 若你训练时用的是CategoricalCrossentropy就替换上面的类
custom_loss = lambda **kwargs: SparseCategoricalCrossentropy(**{k:v for k,v in kwargs.items() if k != 'axis'})

age_model = keras.models.load_model(age_model_path, custom_objects={"SparseCategoricalCrossentropy": custom_loss})
gender_model = keras.models.load_model(gender_model_path, custom_objects={"SparseCategoricalCrossentropy": custom_loss})

• 方案3：将当前环境的TensorFlow版本升级到和训练模型时一致的版本，从根源解决版本兼容问题。

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