加cbam的yolov3模型

CBAM（Convolutional Block Attention Module）是一种用于图像分类和目标检测的注意力机制模块，可以增强模型的表达能力和泛化能力。在 YOLOv3 模型中加入 CBAM 模块可以提高模型的检测精度和减少误检率。

以下是添加 CBAM 模块的 YOLOv3 的代码示例：

import torch.nn as nn
from models.yolo_layers import *

class CBAM(nn.Module):
    def __init__(self, in_channels, reduction_ratio=16):
        super(CBAM, self).__init__()
        self.in_channels = in_channels
        self.mlp = nn.Sequential(
            nn.Linear(in_channels, in_channels // reduction_ratio),
            nn.ReLU(),
            nn.Linear(in_channels // reduction_ratio, in_channels)
        )
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.max_pool = nn.AdaptiveMaxPool2d(1)
        self.conv1 = nn.Conv2d(in_channels=2, out_channels=1, kernel_size=7, stride=1, padding=3, bias=False)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        batch_size, channels, height, width = x.size()
        
        # Channel attention
        avg = self.avg_pool(x).view(batch_size, channels)
        max = self.max_pool(x).view(batch_size, channels)
        out = self.mlp(avg) + self.mlp(max)
        out = self.sigmoid(out).view(batch_size, channels, 1, 1)
        out = x * out
        
        # Spatial attention
        avg = torch.mean(out, dim=1, keepdim=True)
        max = torch.max(out, dim=1, keepdim=True)[0]
        out = torch.cat([avg, max], dim=1)
        out = self.conv1(out)
        out = self.sigmoid(out)
        out = x * out
        
        return out

class YOLOv3CBAM(nn.Module):
    def __init__(self, num_classes=80, anchors=None):
        super(YOLOv3CBAM, self).__init__()
        self.num_classes = num_classes
        self.anchors = anchors

        self.backbone = nn.Sequential(
            CBAMConv2d(3, 32, kernel_size=3, stride=1, padding=1),
            nn.MaxPool2d(kernel_size=2, stride=2),
            CBAMConv2d(32, 64, kernel_size=3, stride=1, padding=1),
            nn.MaxPool2d(kernel_size=2, stride=2),
            CBAMConv2d(64, 128, kernel_size=3, stride=1, padding=1),
            CBAMConv2d(128, 64, kernel_size=1, stride=1, padding=0),
            CBAMConv2d(64, 128, kernel_size=3, stride=1, padding=1),
            nn.MaxPool2d(kernel_size=2, stride=2),
            CBAMConv2d(128, 256, kernel_size=3, stride=1, padding=1),
            CBAMConv2d(256, 128, kernel_size=1, stride=1, padding=0),
            CBAMConv2d(128, 256, kernel_size=3, stride=1, padding=1),
            nn.MaxPool2d(kernel_size=2, stride=2),
            CBAMConv2d(256, 512, kernel_size=3, stride=1, padding=1),
            CBAMConv2d(512, 256, kernel_size=1, stride=1, padding=0),
            CBAMConv2d(256, 512, kernel_size=3, stride=1, padding=1),
            CBAMConv2d(512, 256, kernel_size=1, stride=1, padding=0),
            CBAMConv2d(256, 512, kernel_size=3, stride=1, padding=1),
            nn.MaxPool2d(kernel_size=2, stride=2),
            CBAMConv2d(512, 1024, kernel_size=3, stride=1, padding=1),
            CBAMConv2d(1024, 512, kernel_size=1, stride=1, padding=0),
            CBAMConv2d(512, 1024, kernel_size=3, stride=1, padding=1),
            CBAMConv2d(1024, 512, kernel_size=1, stride=1, padding=0),
            CBAMConv2d(512, 1024, kernel_size=3, stride=1, padding=1),
            CBAMConv2d(1024, 512, kernel_size=1, stride=1, padding=0),
            CBAMConv2d(512, 1024, kernel_size=3, stride=1, padding=1),
        )

        self.head1 = YOLOHead(in_channels=1024, num_anchors=3 * (num_classes + 5))
        self.up1 = nn.Sequential(
            CBAMConv2d(512, 256, kernel_size=1, stride=1, padding=0),
            nn.Upsample(scale_factor=2)
        )
        self.head2 = YOLOHead(in_channels=768, num_anchors=3 * (num_classes + 5))
        self.up2 = nn.Sequential(
            CBAMConv2d(256, 128, kernel_size=1, stride=1, padding=0),
            nn.Upsample(scale_factor=2)
        )
        self.head3 = YOLOHead(in_channels=384, num_anchors=3 * (num_classes + 5))

    def forward(self, x):
        x2, x1, x0 = self.backbone(x)

        out0 = self.head1(x0)
        out1 = self.up1(x0)
        out1 = torch.cat([out1, x1], dim=1)
        out1 = self.head2(out1)
        out2 = self.up2(out1)
        out2 = torch.cat([out2, x2], dim=1)
        out2 = self.head3(out2)

        if self.training:
            return out0, out1, out2
        else:
            return YOLOv3PostProcess(out0, out1, out2, self.anchors, self.num_classes)

其中 `CBAMConv2d` 是使用了 CBAM 的卷积层，其代码如下：

class CBAMConv2d(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride, padding):
        super(CBAMConv2d, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, bias=False)
        self.bn = nn.BatchNorm2d(out_channels)
        self.cbam = CBAM(out_channels)

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

这个 YOLOv3CBAM 模型与原始的 YOLOv3 模型结构相似，不同之处在于加入了 CBAM 模块，可以通过训练来提高检测精度。