如何合并两个结构相同的TensorFlow Keras图像分类模型实现4分类

以下方案可直接将两个结构相同、输入一致的预训练二分类模型合并为端到端四分类模型，仅需传入单张图片即可得到4分类结果，已规避层名重复冲突等常见问题。

前置说明 #

默认两个二分类模型为独立训练的预训练模型，4分类对应两个二分类输出的组合，顺序为：[模型1类别0+模型2类别0, 模型1类别0+模型2类别1, 模型1类别1+模型2类别0, 模型1类别1+模型2类别1]，可根据自身分类定义调整映射规则。

完整实现代码 #

import tensorflow as tf
from tensorflow.keras.layers import Input, Lambda
from tensorflow.keras.models import Model

# 1. 加载预训练模型，修改为你自己的模型存储路径
model1 = tf.keras.models.load_model("your_model1_path.h5")
model2 = tf.keras.models.load_model("your_model2_path.h5")

# 2. 重命名模型、冻结参数避免冲突和参数变更
model1._name = "classifier_1"
model2._name = "classifier_2"
model1.trainable = False
model2.trainable = False

# 3. 定义公共输入层，形状与原模型输入完全一致，灰度图改(50,50,1)，RGBA改(50,50,4)
input_layer = Input(shape=(50, 50, 3), name="common_image_input")

# 4. 同时获取两个子模型的输出
out_1 = model1(input_layer) # 输出形状 (None, 2)
out_2 = model2(input_layer) # 输出形状 (None, 2)

# 5. 计算4分类联合概率，若仅需拼接输出可替换为 tf.concat([out_1, out_2], axis=-1)
def calc_joint_prob(inputs):
    out1, out2 = inputs
    # 广播计算两两组合的联合概率
    out1_exp = tf.expand_dims(out1, axis=-1)
    out2_exp = tf.expand_dims(out2, axis=1)
    joint_probs = out1_exp * out2_exp
    return tf.reshape(joint_probs, (-1, 4))

combined_out = Lambda(calc_joint_prob, name="4class_output_layer")([out_1, out_2])

# 6. 构建合并后的端到端模型
merged_model = Model(inputs=input_layer, outputs=combined_out, name="merged_4class_classifier")

# 验证模型结构
merged_model.summary()

# 7. 保存合并后的模型，后续可直接加载调用
merged_model.save("merged_4class_model.h5")

调用示例 #

# 加载合并后的模型
merged_model = tf.keras.models.load_model("merged_4class_model.h5", custom_objects={"calc_joint_prob": calc_joint_prob})

# 预处理图片为形状 (1, 50, 50, 3) 的张量，和原模型预处理逻辑完全一致
preprocessed_img = your_image_preprocess_func(raw_image)

# 直接推理得到4分类概率，形状为(1,4)
pred_result = merged_model.predict(preprocessed_img, verbose=0)

可选调整项 #

• 若子模型输出为未激活的logits，可在calc_joint_prob中先对out1、out2分别做tf.nn.softmax()再计算联合概率
• 若不需要联合概率，仅需拼接两个二分类的输出，直接将Lambda层替换为tf.concat([out_1, out_2], axis=-1)即可

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