python为现有yolov5模型增加cbam

要为现有的 YOLOv5 模型增加 CBAM，你需要进行以下步骤：

1. 下载 CBAM 模型的代码和权重：https://github.com/Jongchan/attention-module

2. 将 CBAM 模型中的 `cbam.py` 文件复制到 YOLOv5 代码目录中。

3. 在 `models/yolo.py` 文件中，在 `__init__` 函数中导入 `CBAM` 模型：

from cbam import CBAM

4. 在 `models/yolo.py` 文件中，找到 `Conv` 类型的层，并在 `Conv` 层后添加 `CBAM` 模型：

class Conv(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=None, groups=1, bias=True, activate=True, bn=True, name=''):
        super(Conv, self).__init__()
        self.activate = activate
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, groups=groups, bias=bias)
        self.bn = nn.BatchNorm2d(out_channels, eps=1e-3, momentum=0.03) if bn else None
        self.cbam = CBAM(out_channels) # CBAM 模型
        self.name = name

    def forward(self, x):
        x = self.conv(x)
        if self.bn is not None:
            x = self.bn(x)
        if self.activate:
            x = F.relu(x, inplace=True)
        x = self.cbam(x) # 加入 CBAM 模型
        return x

5. 在 `models/yolo.py` 文件中，找到 `C3` 类型的层，并在 `Conv` 层后添加 `CBAM` 模型：

class C3(nn.Module):
    def __init__(self, ch, n=1, shortcut=True, g=1, e=0.5):
        super(C3, self).__init__()
        self.cv1 = Conv(ch, ch // 2, 1, 1)
        self.cv2 = Conv(ch, ch // 2, 1, 1)
        self.cv3 = Conv(ch, ch, 1, 1)
        self.m = nn.Sequential(*[Residual(ch, ch, shortcut, g, e) for _ in range(n)])
        self.cbam = CBAM(ch) # CBAM 模型

    def forward(self, x):
        y = torch.cat([
            self.cv1(x[..., ::2, ::2]),
            self.cv2(x[..., 1::2, ::2]),
            self.cv2(x[..., ::2, 1::2]),
            self.cv3(x[..., 1::2, 1::2])
        ], 1)
        y = self.m(y)
        y = self.cbam(y) # 加入 CBAM 模型
        return y

6. 在 `models/yolo.py` 文件中，找到 `SPP` 类型的层，并在 `Conv` 层后添加 `CBAM` 模型：

class SPP(nn.Module):
    def __init__(self, c1, c2, k=(5, 9, 13)):
        super(SPP, self).__init__()
        c_ = c1 // 2
        self.cv1 = Conv(c1, c_, 1, 1)
        self.cv2 = Conv(c_ * (len(k) + 1), c2, 1, 1)
        self.m = nn.ModuleList([nn.MaxPool2d(kernel_size=x, stride=1, padding=x // 2) for x in k])
        self.cbam = CBAM(c2) # CBAM 模型

    def forward(self, x):
        x = self.cv1(x)
        x = torch.cat([x] + [m(x) for m in self.m], 1)
        x = self.cv2(x)
        x = self.cbam(x) # 加入 CBAM 模型
        return x

7. 在 `models/yolo.py` 文件中，找到 `Detect` 类型的层，并在 `Conv` 层后添加 `CBAM` 模型：

class Detect(nn.Module):
    stride = None  # strides computed during build
    export = False  # onnx export

    def __init__(self, nc, anchors, ch=256):
        super(Detect, self).__init__()
        self.nc = nc  # number of classes
        self.no = nc + 5  # number of outputs per anchor
        self.na = len(anchors)  # number of anchors
        self.anchors = torch.Tensor(anchors)
        self.register_buffer('anchors_scaled', self.anchors.clone().view(1, -1, 1, 1, 2) / ch)
        self.m = nn.ModuleList(nn.Conv2d(x, self.no * self.na, 1) for x in [ch * 4, ch * 8, ch * 16])
        self.cbam = CBAM(ch * 16) # CBAM 模型

    def forward(self, x):
        z = []
        for i in range(3):
            xi = x[i]
            mi = self.m[i](xi)
            zi = mi.reshape(mi.shape[0], self.na, self.no, mi.shape[2], mi.shape[3]).permute(0, 1, 3, 4, 2).contiguous()
            if not self.training:
                if self.export:
                    return zi
                else:
                    xi_shape = xi.shape[-2:]
                    if self.stride is None:
                        self.stride = [x.shape[-1] / xi_shape[-1] for x in x]
                    self.stride[i] = x[0].shape[-1] / xi_shape[-1]
                    self.anchors /= self.stride[i]
                    self.anchors_scaled = self.anchors.clone().view(1, -1, 1, 1, 2) / xi_shape[::-1]
            z.append(zi)
        x = torch.cat(z, 1)
        x = self.cbam(x) # 加入 CBAM 模型
        return x

8. 最后运行你的 YOLOv5 代码，并检查是否已经加入了 CBAM 模型。