Keras迁移学习加载U-Net权重时Layer count mismatch错误排查

核心问题原因 #

你将预训练U-Net作为单个整体层嵌入新模型，导致新模型的层结构（仅4层）与预训练权重的层结构（23层）完全不匹配，Keras无法对应分配权重。

解决方案一：拆解预训练模型层构建新模型（推荐） #

这是最稳妥的方式，直接复用预训练模型的所有层结构和权重：

1. 加载完整的预训练U-Net模型

   from keras.models import load_model
   
   # 加载你之前训练好的完整U-Net模型
   pretrained_unet = load_model('你的预训练U-Net.h5')
   

2. （可选）冻结预训练层（迁移学习初期避免破坏已训练权重）

   for layer in pretrained_unet.layers:
       layer.trainable = False
   # 后续训练可根据需求解冻部分卷积块，比如解冻最后1-2个卷积块
   

3. 构建新模型的输入与转通道层

   from keras.layers import Input, Conv2D
   from keras.models import Model
   
   # 新数据集的5通道输入
   new_input = Input(shape=(128, 128, 5))
   # 新增卷积层将5通道转为3通道，适配预训练U-Net的输入
   # 注意padding设为same保证输出尺寸和输入一致，激活函数按需选择
   conv_to_3ch = Conv2D(3, (3, 3), padding='same', activation='relu')(new_input)
   

4. 拼接预训练层与输出层

   # 将转通道后的特征图传入预训练U-Net
   unet_features = pretrained_unet(conv_to_3ch)
   # 新增语义分割的输出层（注意：语义分割建议用Conv2D而非Dense，避免丢失空间信息）
   # 示例：多分类用softmax，二分类用sigmoid，num_classes替换为你的类别数
   final_output = Conv2D(num_classes, (1, 1), activation='softmax')(unet_features)
   # 构建完整的迁移学习模型
   new_model = Model(inputs=new_input, outputs=final_output)
   

5. 验证模型结构
   运行new_model.summary()，你会看到预训练的23层全部被拆解显示，加上新增的转通道层和输出层，此时模型权重已自动继承预训练U-Net的参数，无需额外加载权重文件。

解决方案二：按层名称匹配加载权重（适用于仅保存权重文件的场景） #

如果你之前只保存了预训练模型的权重文件（而非完整模型），可以通过层名称匹配来加载：

1. 严格复刻预训练U-Net的层结构，同时新增转通道层和输出层，确保预训练部分的层名称与原模型完全一致
2. 加载权重时指定by_name=True，Keras会忽略新增的层，仅匹配名称相同的层加载权重：

   new_model.load_weights('预训练权重文件.h5', by_name=True)
   

注意事项 #

• 语义分割任务中，不要用Dense层作为输出，Dense会将二维特征图展平为一维，丢失空间位置信息，建议用Conv2D(1, (1,1), activation='sigmoid')（二分类）或Conv2D(num_classes, (1,1), activation='softmax')（多分类）。
• 迁移学习后期，可逐步解冻预训练U-Net的上层卷积块，联合微调，提升模型在新数据集上的性能。

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