使用YOLOv7预训练分割权重(yolov7-mask.pt)遇AttributeError求助

问题分析 #

你遇到的AttributeError: Can't get attribute 'Merge' on <module 'models.common' from '/content/yolov7/models/common.py'>错误，本质是yolov7-mask.pt预训练权重对应的模型结构中包含Merge模块，但当前使用的YOLOv7代码库的models/common.py里没有定义这个类。普通检测用的yolov7.pt不需要这个模块，所以能正常加载，而分割分支的权重依赖这个自定义模块。

可行解决方案 #

1. 切换到官方YOLOv7分割专用代码分支 #

YOLOv7的分割功能是在特定版本的代码中实现的，普通检测版代码不包含分割所需的模块。

• 确保使用的是官方支持分割任务的YOLOv7代码，此时models/common.py中会自带Merge类。
• 重新用yolov7-mask.pt启动训练，模型结构与权重匹配后即可正常加载。

2. 手动在models/common.py中添加Merge模块 #

如果不想替换整个代码库，直接在当前的models/common.py末尾添加以下代码：

class Merge(nn.Module):
    def __init__(self, c1, c2):
        super().__init__()
        self.conv = Conv(c1, c2, 1, 1)
        self.bn = nn.BatchNorm2d(c2)
        self.act = nn.SiLU()

    def forward(self, x):
        return self.act(self.bn(self.conv(x)))

保存文件后重新运行训练脚本，就能正常加载yolov7-mask.pt权重。

3. 调整权重加载的严格模式（临时 workaround） #

如果只是想跳过模块不匹配的问题（不推荐，可能丢失部分预训练效果），可以修改模型加载相关代码：
在加载权重的model.load_state_dict步骤中添加strict=False参数，示例：

model.load_state_dict(torch.load(weights, map_location=device)['model'], strict=False)

这种方法会忽略不匹配的模块，但预训练权重的作用会打折扣，优先推荐前两种方案。

额外建议 #

• 确认自定义数据集的掩码标注格式符合YOLOv7分割的要求，避免后续训练出现格式类错误。
• 预训练权重能大幅缩短小数据集的训练周期，优先解决权重加载问题而非从零训练。

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