pytorch中加入注意力机制（CBAM），以yolov5为例

书本下载地址：https://www.taylorfrancis.com/books/matlab-simulations-radar-systems-design-bassem-mahafza-atef-elsherbeni/10.1201/9780203502556

YY_oot

34663人浏览 · 2021-06-14 14:25:18

YY_oot · 2021-06-14 14:25:18 发布

代码实现：CBAM.PyTorch
参考文章：pytorch中加入注意力机制（CBAM），以ResNet为例
yolo理论解读：yolov5l.yaml ； yolo.py

主要改 .yaml, yolo.py, commom.py 这三个文件

① .yaml
yolov5提供了s、m、l、x四种，CBAM要加在backbone的位置。
以我的代码为例，我在第一个卷积后加入了CBAM，把相应的 ”Conv“ 改为了”Conv_CBAM“

backbone:
  [[-1, 1, Focus, [64, 3]],  # 0-P1/2
   [-1, 1, Conv_CBAM, [128, 3, 2]],  # 1-P2/4 # add CBAM
   [-1, 3, BottleneckCSP, [128]],
   [-1, 1, Conv, [256, 3, 2]],  # 3-P3/8
   [-1, 9, BottleneckCSP, [256]],
   [-1, 1, Conv, [512, 3, 2]],  # 5-P4/16
   [-1, 9, BottleneckCSP, [512]],
   [-1, 1, Conv, [1024, 3, 2]],  # 7-P5/32
   [-1, 1, SPP, [1024, [5, 9, 13]]],
   [-1, 3, BottleneckCSP, [1024, False]],  # 9
  ]

② yolo.py
模型文件中的parse_model函数，用来读入模型yaml中的参数定义，代码更改如下：
【原】

from models.common import Conv, Bottleneck, SPP, DWConv, Focus, BottleneckCSP, Concat, NMS

【改】

from models.common import Conv, Bottleneck, SPP, DWConv, Focus, BottleneckCSP, Concat, NMS, Conv_CBAM

【原】

if m in [Conv, Bottleneck, SPP, DWConv, MixConv2d, Focus, CrossConv, BottleneckCSP, C3]:

【改】

if m in [Conv, Bottleneck, SPP, DWConv, MixConv2d, Focus, CrossConv, BottleneckCSP, C3, Conv_CBAM]:

③ commom.py
yolo.py会调用commom.py里的函数，该部分是backbone各个模块参数讲解
我们加入以下代码：
【标准卷积层 + CBAM 模块】

# 标准卷积层 + CBAM
class Conv_CBAM(nn.Module):
    # Standard convolution
    def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True):  # ch_in, ch_out, kernel, stride, padding, groups
        super(Conv_CBAM, self).__init__()
        self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p), groups=g, bias=False)
        self.bn = nn.BatchNorm2d(c2)
        self.act = nn.Hardswish() if act else nn.Identity()
        self.ca = ChannelAttention(c2)
        self.sa = SpatialAttention()

    def forward(self, x):
        x = self.act(self.bn(self.conv(x)))
        x = self.ca(x) * x
        x = self.sa(x) * x
        return x

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

【CBAM定义模块】

# add CBAM
class ChannelAttention(nn.Module):
    def __init__(self, in_planes, ratio=16):
        super(ChannelAttention, self).__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.max_pool = nn.AdaptiveMaxPool2d(1)

        self.fc1   = nn.Conv2d(in_planes, in_planes // 16, 1, bias=False)
        self.relu1 = nn.ReLU()
        self.fc2   = nn.Conv2d(in_planes // 16, in_planes, 1, bias=False)

        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        avg_out = self.fc2(self.relu1(self.fc1(self.avg_pool(x))))
        max_out = self.fc2(self.relu1(self.fc1(self.max_pool(x))))
        out = avg_out + max_out
        return self.sigmoid(out)


class SpatialAttention(nn.Module):
    def __init__(self, kernel_size=7):
        super(SpatialAttention, self).__init__()

        assert kernel_size in (3, 7), 'kernel size must be 3 or 7'
        padding = 3 if kernel_size == 7 else 1

        self.conv1 = nn.Conv2d(2, 1, kernel_size, padding=padding, bias=False)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        avg_out = torch.mean(x, dim=1, keepdim=True)
        max_out, _ = torch.max(x, dim=1, keepdim=True)
        x = torch.cat([avg_out, max_out], dim=1)
        x = self.conv1(x)
        return self.sigmoid(x)