python为现有yolov5模型增加cbam

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

1. 导入必要的库和模块

import torch
from torch import nn
from yolov5.models.yolo import YOLOv5

2. 定义 CBAM 模块

class CBAM(nn.Module):
    def __init__(self, channels, reduction=16):
        super().__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.max_pool = nn.AdaptiveMaxPool2d(1)

        self.fc1 = nn.Conv2d(channels, channels // reduction, kernel_size=1, bias=False)
        self.relu = nn.ReLU(inplace=True)
        self.fc2 = nn.Conv2d(channels // reduction, channels, kernel_size=1, bias=False)
        self.sigmoid_channel = nn.Sigmoid()

        self.conv_after_concat = nn.Conv2d(2, 1, kernel_size=7, stride=1, padding=3, bias=False)
        self.sigmoid_spatial = nn.Sigmoid()

    def forward(self, x):
        avg_out = self.fc2(self.relu(self.fc1(self.avg_pool(x))))
        max_out = self.fc2(self.relu(self.fc1(self.max_pool(x))))
        channel_out = self.sigmoid_channel(avg_out + max_out)

        spatial_out = torch.cat([avg_out, max_out], dim=1)
        spatial_out = self.conv_after_concat(spatial_out)
        spatial_out = self.sigmoid_spatial(spatial_out)

        return x * channel_out * spatial_out

3. 修改 YOLOv5 模型

在 YOLOv5 模型的定义中，我们需要找到每个卷积层并添加 CBAM 模块。这里以 v5s 为例，代码如下：

class YOLOv5s(YOLOv5):
    def __init__(self, num_classes=80):
        super().__init__(ch=32, in_channels=3, num_classes=num_classes)

        # modify yolov5s to add cbam modules
        for m in self.model.children():
            if isinstance(m, nn.Conv2d):
                m.register_forward_pre_hook(self.hook_fn_cbam)

    def hook_fn_cbam(self, module, input):
        input[0] = CBAM(input[0].shape[1])(input[0])

在这个例子中，我们遍历了 YOLOv5s 模型的每个子模块，并找到了所有的卷积层。 然后，我们为每个卷积层添加了一个 forward_pre_hook，这个 hook 将在模型前向传递过程中被调用。hook_fn_cbam 函数中，我们使用 CBAM 模块对输入进行处理，并将结果作为新的输入。

4. 实例化模型

model = YOLOv5s(num_classes=80)

现在，你已经成功地为现有的 YOLOv5 模型增加了 CBAM 模块。接下来，你可以使用该模型进行训练和测试。