YOLOv5同类别训练模型合并失败：参数尺寸不匹配问题求助

问题原因分析 #

你遇到的参数尺寸不匹配，核心原因是输出层维度和类别数直接绑定：

• YOLOv5输出层的维度计算公式为 (类别数 + 5) * 3，其中5对应坐标(x,y)、宽高(w,h)、置信度(confidence)，3是每个特征层的锚框数量。
• 你的两个训练模型针对4类任务，所以输出层维度是 (4+5)*3=27，对应报错里的[27,128,1,1]；而你直接加载的原始yolov5s.pt默认是80类，输出层维度是(80+5)*3=255，也就是报错里的[255,128,1,1]，两者维度自然不匹配。

解决方案 #

步骤1：创建匹配类别数的新模型 #

合并权重前，必须先基于yolov5s结构创建4类的新模型，而非直接使用原始80类的模型结构。

步骤2：修改合并代码（merge.py） #

调整后的核心代码逻辑如下：

import torch
from models.yolo import Model
from utils.general import check_yaml

# 加载YOLOv5s的配置文件
cfg = check_yaml('models/yolov5s.yaml')
# 创建4类的新模型，nc参数指定类别数
new_model = Model(cfg, ch=3, nc=4)

# 加载训练好的ModelA和ModelB
model_a = torch.load('ModelA.pt')['model'].float()
model_b = torch.load('ModelB.pt')['model'].float()

# 遍历参数，对两个模型的对应参数取平均后赋值给新模型
for (name_a, param_a), (name_b, param_b), (name_new, param_new) in zip(model_a.named_parameters(), model_b.named_parameters(), new_model.named_parameters()):
    assert name_a == name_b == name_new, "参数名称不匹配，模型结构可能不一致"
    param_new.data = (param_a.data + param_b.data) / 2.0

# 保存合并后的模型
torch.save({'model': new_model.half()}, 'merged_model.pt')

额外说明 #

训练时出现的Seaborn警告和plot_images的AttributeError属于训练过程中的非致命问题，不影响模型权重的有效性，可忽略或通过升级seaborn、matplotlib库解决，但和本次合并问题无关。

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