Keras Functional API构建ResNet权重保存异常及推理结果差问题求助

你的问题核心是Graph模式下模型权重保存/加载后推理异常，但Eager模式正常，根源集中在模型构建的无状态性、多GPU权重保存逻辑、BatchNorm状态处理这几个点上，以下是具体修复步骤：

1. 移除Bottleneck函数中的全局变量identitys #

全局变量是Graph模式的大忌，会破坏模型结构的一致性，导致保存的权重与加载后的模型层不匹配。修改Bottleneck函数为局部变量实现：

def Bottleneck(inputs, out_channel, name, downsample, strides=1):
    expansion = 4
    key = out_channel * expansion
    identity = inputs

    # 替换全局变量为局部变量
    if downsample:
        identity_out = layers.Conv2D(key, kernel_size=1, strides=strides,
                                  use_bias=False, kernel_initializer='he_normal',
                                  padding="SAME", kernel_regularizer=l2(1.e-5), name=name + "ds_conv")(identity)
        identity_out = layers.BatchNormalization(momentum=0.9, epsilon=1e-5, name=name + "ds_normal")(identity_out)
    else:
        identity_out = inputs

    xb = layers.Conv2D(out_channel, kernel_size=1, use_bias=False, kernel_initializer='he_normal',
                       kernel_regularizer=l2(1.e-4), name=name + "Conv2D_1")(inputs)
    xb = layers.BatchNormalization(momentum=0.9,
                                   epsilon=1e-5, name=name + "BN_1")(xb)
    xb = layers.Activation(tf.keras.activations.swish, name=name + "ACT_1")(xb)

    xb = layers.Conv2D(out_channel, kernel_size=3, use_bias=False, strides=strides, padding="SAME",
                       kernel_initializer='he_normal', kernel_regularizer=l2(1.e-4), name=name + "Conv2D_2")(xb)
    xb = layers.BatchNormalization(momentum=0.9,
                                   epsilon=1e-5, name=name + "BN_3")(xb)
    xb = layers.ReLU(name=name + "ReLU")(xb)

    xb = layers.Conv2D(key, kernel_size=1, use_bias=False,
                       kernel_initializer='he_normal',
                       kernel_regularizer=l2(1.e-4),
                       name=name + "Conv2D_3")(xb)

    xb = layers.BatchNormalization(momentum=0.9,
                                   epsilon=1e-5,
                                   name=name + "BN_4")(xb)

    xb = layers.Add(name=name + "addition")([identity_out, xb])
    xb = layers.BatchNormalization(momentum=0.9,
                                   epsilon=1e-5,
                                   name=name + "Last_BN")(xb)
    xb = layers.ReLU(name=name + "LastReLU")(xb)

    return xb

2. 修正多GPU训练的权重保存逻辑 #

分布式训练中，必须保存单GPU原始模型的权重，而非分布式包装后的模型。修改训练代码：

# 先创建单GPU基础模型
base_model = ResnetBuilder.resnet50(448, 448, 5)
# 在策略范围内包装为分布式模型
with strategy.scope():
    model = tf.keras.models.clone_model(base_model)
    model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])

checkpoint_path = save_train_data + "/model.{epoch:02d}.h5"
# 回调函数传入单GPU的base_model
save_weight = MultiGPUCheckpointCallback(filepath=checkpoint_path,
                                         base_model=base_model,
                                         save_weights_only=True)

history = model.fit(x=train_dataset, validation_data=valid_dataset,
                    steps_per_epoch=int(np.ceil(training_step_nums / BATCH_SIZE)),
                    validation_steps=int(np.ceil(validation_nums / BATCH_SIZE)),
                    epochs=EPOCHES, verbose="auto", callbacks=[save_weight])

3. 推理时强制切换模型到评估模式 #

加载权重后，必须让模型进入评估状态，确保BatchNorm等层使用训练好的移动均值/方差，而非继续更新：

def load_trained_model():
    model = ResnetBuilder.resnet50(448, 448, 5)
    model_name = r"./model.30.h5"
    # 移除by_name=True，确保所有权重精确匹配加载
    model.load_weights(model_name)
    # 切换到评估模式
    model.trainable = False
    model.compile()  # 确保模型状态初始化完成
    print('model load success.')
    return model

4. 简化推理代码，避免不必要的模型包装 #

推理时无需重新创建heatmap_model，直接使用原模型即可：

def get_category_name(full_image_path, model):
    img = preprocess_image(full_image_path)
    img_tensor = np.expand_dims(img, axis=0)

    # 直接用原模型预测
    predictions = model.predict(img_tensor, verbose=0)
    category_id = np.argmax(predictions[0])
    label_name = ['A1', 'A2', 'A3', "A4", "A5"]
    category_name = label_name[category_id]

    return category_name

关键原理说明 #

• 全局变量破坏Graph模式：Graph模式下Keras构建静态计算图，全局变量会导致模型层的引用混乱，保存的权重与加载后的模型结构不匹配；而Eager模式是动态执行，不会触发该问题。
• BatchNorm状态切换：训练时BatchNorm处于更新均值方差的模式，推理时必须切换到评估模式，否则会用实时计算的均值方差导致结果异常。
• 多GPU权重一致性：分布式模型的权重分散在多个GPU上，只有保存单GPU原始模型的权重，才能在单GPU环境下正确加载复用。

内容的提问来源于stack exchange，提问作者张家瑜