YOLOv5替换Backbone为ResNet-50时的形状匹配错误

问题背景 #

我是一名研究者，想要将YOLOv5的Backbone从Darknet替换为ResNet。已在commons.py中添加MaxPooling2D和ResBlock两个类，并在Yolo.py中注册类名。模型摘要可成功打印，但运行时出现形状匹配错误，相关代码及报错如下：

相关代码 #

class Conv(nn.Module):
    # Standard convolution with args(ch_in, ch_out, kernel, stride, padding, groups, dilation, activation)
    default_act = nn.SiLU()  # default activation

    def __init__(self, c1, c2, k=1, s=1, p=None, g=1, d=1, act=True):
        super().__init__()
        self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p, d), groups=g, dilation=d, bias=False)
        self.bn = nn.BatchNorm2d(c2)
        self.act = self.default_act if act is True else act if isinstance(act, nn.Module) else nn.Identity()

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

    def forward_fuse(self, x):
        return self.act(self.conv(x))

class MaxPooling2D(nn.Module):
    # MaxPooling2D layer with args(kernel, stride, padding)
    def __init__(self, k=2, s=2, p=0):
        super().__init__()
        self.maxpool = nn.MaxPool2d(k, s, p)

    def forward(self, x):
        return self.maxpool(x)

class ResBlock(nn.Module):
    def __init__(self, c1, num_repeats):
        super().__init__()
        self.blocks = nn.Sequential(*[
            nn.Identity() if i == 0 else Conv(c1, c1, k=3, s=1, act=True) for i in range(num_repeats)
        ])

    def forward(self, x):
        return torch.cat(x, self.blocks(x))

# Backbone配置
backbone: [
    [-1, 1, "Conv", [64, 6, 2, 2]],  # 0 - Conv1
    [-1, 1, "MaxPooling2D", [3, 2]],   # 1 - MaxPool
    [-1, 3, "ResBlock", [64]],         # 2 - Res2a, Res2b, Res2c
    [-1, 1, "Conv", [128, 3, 2]], # 3 - Conv3
    [-1, 4, "ResBlock", [128]],        # 4 - Res3a, Res3b, Res3c, Res3d
    [-1, 1, "Conv", [256, 3, 2]], # 5 - Conv4
    [-1, 6, "ResBlock", [256]],        # 6 - Res4a, Res4b, Res4c, Res4d, Res4e, Res4f
    [-1, 1, "Conv", [512, 3, 2]], # 7 - Conv5
    [-1, 3, "ResBlock", [512]],        # 8 - Res5a, Res5b, Res5c
    [-1, 1, "SPPF", [1024, 5]]         # 9 - SPPF
]

报错信息 #

Traceback (most recent call last):
  File "train.py", line 647, in <module>
    main(opt)
  File "train.py", line 536, in main
    train(opt.hyp, opt, device, callbacks)
  File "train.py", line 130, in train
    model = Model(cfg or ckpt['model'].yaml, ch=3, nc=nc, anchors=hyp.get('anchors')).to(device)  # create
  File "/home/dev/Documents/yolov5/models/yolo.py", line 195, in __init__
    m.stride = torch.tensor([s / x.shape[-2] for x in forward(torch.zeros(1, ch, s, s))])  # forward
  File "/home/dev/Documents/yolov5/models/yolo.py", line 194, in <lambda>
    forward = lambda x: self.forward(x)[0] if isinstance(m, Segment) else self.forward(x)
  File "/home/dev/Documents/yolov5/models/yolo.py", line 209, in forward
    return self._forward_once(x, profile, visualize)  # single-scale inference, train
  File "/home/dev/Documents/yolov5/models/yolo.py", line 121, in _forward_once
    x = m(x)  # run
  File  "/home/dev/Documents/yolov5/models/common.py", line 90, in forward
    return torch.cat(x, self.blocks(x))
  File "/home/dev/Documents/yolov5/models/common.py", line 68, in forward
    return self.act(self.bn(self.conv(x)))
  File "/home/dev/.cache/pypoetry/virtualenvs/yolov5-FT1Hnn5N-py3.8/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/dev/.cache/pypoetry/virtualenvs/yolov5-FT1Hnn5N-py3.8/lib/python3.8/site-packages/torch/nn/modules/conv.py", line 463, in forward
    return self._conv_forward(input, self.weight, self.bias)
 RuntimeError: Given groups=1, weight of size [8, 8, 3, 3], expected input[1, 32, 13, 13] to have 8 channels, but got 32 channels instead

问题分析与修复方案 #

1. ResBlock forward方法逻辑错误 #

原代码中torch.cat(x, self.blocks(x))存在两处问题：

• torch.cat的正确用法是接收张量列表，而非两个单独张量
• ResNet残差块的核心是shortcut相加，不是张量拼接

修复后的forward方法：

def forward(self, x):
    return x + self.blocks(x)

2. ResBlock的模块构建逻辑错误 #

原代码中通过nn.Identity() if i ==0 else Conv(...)构建的模块完全不符合ResNet残差块结构，导致通道数不匹配。标准ResNet BasicBlock的结构应为两次3x3卷积（通道数一致时），或1x1降维→3x3→1x1升维（通道数变化时）。

修复后的ResBlock实现（适配YOLOv5风格）：

class ResBlock(nn.Module):
    def __init__(self, c1, c2, shortcut=True, g=1, e=0.5):
        # c1: 输入通道数, c2: 输出通道数, shortcut: 是否启用残差连接
        super().__init__()
        c_ = int(c2 * e)  # 中间层通道数
        self.cv1 = Conv(c1, c_, 3, 1)
        self.cv2 = Conv(c_, c2, 3, 1, g=g)
        self.add = shortcut and c1 == c2  # 仅当输入输出通道一致时启用残差

    def forward(self, x):
        return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x))

3. Backbone配置同步调整 #

对应修改后的ResBlock，需要更新backbone配置的参数（新增输出通道数参数）：

backbone: [
    [-1, 1, "Conv", [64, 6, 2, 2]],  # 0 - Conv1
    [-1, 1, "MaxPooling2D", [3, 2]],   # 1 - MaxPool
    [-1, 3, "ResBlock", [64, 64]],     # 2 - 3个BasicBlock，输入输出均为64通道
    [-1, 1, "Conv", [128, 3, 2]], # 3 - Conv3，升维至128
    [-1, 4, "ResBlock", [128, 128]],   # 4 - 4个BasicBlock
    [-1, 1, "Conv", [256, 3, 2]], # 5 - Conv4，升维至256
    [-1, 6, "ResBlock", [256, 256]],   # 6 - 6个BasicBlock
    [-1, 1, "Conv", [512, 3, 2]], # 7 - Conv5，升维至512
    [-1, 3, "ResBlock", [512, 512]],   # 8 - 3个BasicBlock
    [-1, 1, "SPPF", [1024, 5]]         # 9 - SPPF
]

4. 额外排查建议 #

• 运行前打印每一层的输出形状，确认通道数和特征图尺寸是否符合预期
• 确保MaxPooling2D的参数与原Darknet中的MaxPool一致，避免尺寸不匹配

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