yolov5添加注意力机制（CBAM）

Python资料网推荐yolov5添加注意力机制（CBAM）这篇文章给大家，欢迎收藏Python资料网享受知识的乐趣

CBAM注意力机制

CBAM注意力机制是一种用于计算机视觉任务的注意力机制，它通过对图像的通道和空间维度进行注意力加权来提升模型的性能。它由两个关键组件组成：通道注意力和空间注意力。

通道注意力模块主要用于对输入的通道特征图进行加权。它通过使用全局平均池化来获取全局的通道特征表示，然后经过一层全连接网络，生成通道注意力权重。这些权重用于调整每个通道的重要性，使得模型能够自动学习到不同通道之间的相关性和重要性，从而更好地利用通道信息。

空间注意力模块主要用于对输入的空间特征图进行加权。它通过使用逐通道的最大池化和平均池化来提取不同尺度的空间信息，然后经过一层全连接网络，生成空间注意力权重。这些权重用于调整每个空间位置的重要性，使得模型能够自适应地关注不同位置的细节信息，从而提升模型对空间变化的感知能力。

通过结合通道注意力和空间注意力，CBAM注意力机制能够自适应地调整不同通道和空间位置的重要性，从而提升模型对图像特征的表达能力。这使得模型能够更好地捕捉图像中的关键信息，提高任务的准确性和鲁棒性。

引入CBAM后，YOLOv5在基础骨干网的每个残差单元的最后添加了CBAM模块，并在中间层的SPP（Spatial Pyramid Pooling）层后也添加了CBAM模块。这个修改使得模型能够更好地捕捉特征之间的依赖关系，并提高特征表示能力和检测精度。此外，YOLOv5的检测头也添加了CBAM模块，从而使模型能够自适应地学习不同尺度特征之间的依赖关系，进一步提高检测精度。总之，通过引入CBAM注意力机制，YOLOv5改进了网络的基础骨干网、中间层和检测头，使得模型能够更好地捕捉特征之间的依赖关系，提高特征表示能力和检测精度。

CBAM注意力机制代码

步骤1 

在common.py中添加CBAM模块

将下面的CBAM模块的代码复制粘贴到common.py文件的末尾。

# 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.f1 = nn.Conv2d(in_planes, in_planes // ratio, 1, bias=False)
        self.relu = nn.ReLU()
        self.f2 = nn.Conv2d(in_planes // ratio, in_planes, 1, bias=False)
        self.sigmoid = nn.Sigmoid()
    def forward(self, x):
        avg_out = self.f2(self.relu(self.f1(self.avg_pool(x))))
        max_out = self.f2(self.relu(self.f1(self.max_pool(x))))
        out = self.sigmoid(avg_out + max_out)
        return 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
        # (特征图的大小-算子的size+2*padding)/步长+1
        self.conv = nn.Conv2d(2, 1, kernel_size, padding=padding, bias=False)
        self.sigmoid = nn.Sigmoid()
    def forward(self, x):
        # 1*h*w
        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)
        #2*h*w
        x = self.conv(x)
        #1*h*w
        return self.sigmoid(x)
    
class CBAM(nn.Module):
    def __init__(self, c1, c2, ratio=16, kernel_size=7):  # ch_in, ch_out, number, shortcut, groups, expansion
        super(CBAM, self).__init__()
        self.channel_attention = ChannelAttention(c1, ratio)
        self.spatial_attention = SpatialAttention(kernel_size)
    def forward(self, x):
        out = self.channel_attention(x) * x
        # c*h*w
        # c*h*w * 1*h*w
        out = self.spatial_attention(out) * out
        return out

步骤2

在yolo.py文件中加入类名

首先在yolo.py文件中找到parse_model函数，然后将CBAM添加到这个注册表里

步骤3

创建自定义yaml文件

在models文件夹中复制yolov5s.yaml，粘贴并命名为yolov5s_CBAM.yaml

步骤4

修改yolov5s_CBAM.yaml文件 

本步骤是修改yolov5s_CBAM.yaml，将CBAM添加到我们想添加的位置。在这里，我将[-1，1，CBAM，[1024]]添加到SPPF的上一层，即下图中所示位置。

注意：注意力机制可以加在Backbone、Neck、Head等部分，常见的有两种：一种是在主干的SPPF前面添加一层；二是将Backbone中的C3全部替换。不同的位置效果可能不同，需要我们去反复测试。

步骤5 

验证是否添加成功

在yolo.py文件里，将配置改为我们刚才自定义的yolov5s_CBAM.yaml

然后运行yolo.py，得到结果

步骤6

修改train.py中的'--cfg'默认参数

 在train.py文件中找到 parse_opt函数，然后将第二行'--cfg'的default改为 'models/yolov5s_CBAM.yaml'，然后就可以开始进行训练了

实验环境

实验运行的硬件加速器为Tesla T4（16G）显卡，torch版本为1.13

实验结果

总体来说，加入CBAM注意力机制的作用还是比较明显的，但注意力机制也不是适合所有的被检测物，还需进行探索。和未改进的检测效果有提升，为改进的实验结果图我就不放了。

希望大家能尽快找到适合的改进方案。

版权声明：本文来自互联网用户投稿，该文观点仅代表作者本人，不代表本站立场。本站仅提供信息存储空间服务，不拥有所有权，不承担相关法律责任。如若内容造成侵权/违法违规/事实不符，请联系邮箱：jacktools123@163.com进行投诉反馈，一经查实，立即删除！