加载已保存Keras模型失败：Tidx属性类型不匹配错误

错误详情 #

ValueError: Inconsistent values for attr 'Tidx' DT_FLOAT vs. DT_INT32 while building NodeDef

涉及节点：tf_op_layer_Mean_17/Mean_17，对应Op定义：

name=Mean; signature=input:T, reduction_indices:Tidx -> output:T; attr=keep_dims:bool,default=false; attr=T:type,allowed=[DT_FLOAT, DT_DOUBLE, DT_INT32, DT_UINT8, DT_INT16, ..., DT_UINT16, DT_COMPLEX128, DT_HALF, DT_UINT32, DT_UINT64]; attr=Tidx:type,default=DT_INT32,allowed=[DT_INT32, DT_INT64]

复现代码 #

import tensorflow as tf
import numpy as np
from tensorflow.keras import Input, Model

tf.compat.v1.disable_eager_execution()
# tf.compat.v1.enable_eager_execution()

inputs = Input(shape=(2,))
output_loss = tf.keras.backend.mean(inputs)
outputs = [inputs, output_loss]
model = Model(inputs, outputs)

loss = tf.reduce_mean(output_loss)  # 触发加载错误
# loss = tf.math.rsqrt(output_loss)  # 无错误
model.add_loss(loss)

model.compile(optimizer="adam", loss=[None] * len(model.outputs))
model.fit(np.random.random((5, 2)), epochs=2)

model.save("my_model_.h5")

# 加载模型时报错
model_ = tf.keras.models.load_model("my_model_.h5", compile=False)
model_.summary()

问题分析 #

问题出在使用tf.reduce_mean作为自定义损失时，模型保存后加载阶段，TensorFlow静态图构建时对reduction_indices（即Tidx参数）的类型推断出现错误——将其识别为DT_FLOAT类型，而该Op要求Tidx必须是DT_INT32或DT_INT64。改用tf.math.rsqrt无问题，是因为该操作不涉及Tidx类型参数。

解决建议 #

1. 移除多余的tf.reduce_mean #

output_loss已经是tf.keras.backend.mean(inputs)计算出的标量，再对其调用tf.reduce_mean完全冗余，直接用output_loss作为损失即可：

loss = output_loss  # 替代原loss定义

2. 显式指定Tidx类型 #

如果确实需要对非标量张量调用tf.reduce_mean，显式指定轴参数的类型为int32/int64，避免类型推断错误：

# 假设需要对轴0求均值，强制转换轴类型
loss = tf.reduce_mean(output_loss, axis=tf.cast(0, tf.int32))

3. 统一使用Keras后端API #

用tf.keras.backend.mean替代tf.reduce_mean，保持与模型构建时的API一致，减少类型推断冲突：

loss = tf.keras.backend.mean(output_loss)

4. 启用Eager Execution #

注释掉tf.compat.v1.disable_eager_execution()，启用Eager模式，动态图构建逻辑可规避静态图中的类型匹配问题：

# tf.compat.v1.disable_eager_execution()
tf.compat.v1.enable_eager_execution()

5. 正确封装自定义损失层 #

如果需要自定义损失逻辑，确保在Layer中显式处理参数类型：

class CustomMeanLoss(tf.keras.layers.Layer):
    def call(self, inputs):
        # 显式指定轴类型为int32
        return tf.reduce_mean(inputs, axis=tf.constant(0, dtype=tf.int32))

# 使用自定义层定义损失
loss_layer = CustomMeanLoss()
loss = loss_layer(output_loss)
model.add_loss(loss)

内容的提问来源于stack exchange，提问作者Mihai.Mehe