YOLOv5添加注意力机制的具体步骤

最新推荐文章于 2025-11-08 21:10:17 发布

原创

最新推荐文章于 2025-11-08 21:10:17 发布 · 2.6w 阅读

860 ·

CC 4.0 BY-SA版权

文章标签：

#深度学习 #pytorch #python

该文详细介绍了如何在YOLOv5-5.0模型中集成CBAM和SE注意力机制。首先在common.py文件中添加CBAMC3和SE模块的代码，接着在yolo.py中插入相应的判断语句，最后修改yaml配置文件以应用这些模块。通过这些步骤，CBAM和SE被成功添加到模型的主干网络中，可用于提升模型的特征学习能力。

本文以CBAM和SE注意力机制的添加过程为例，主要介绍了向YOLOv5中添加注意力机制的具体步骤

本文在此篇博客的基础上向YOLOv5-5.0版本代码中添加注意力机制
yolov5模型训练———使用yolov5训练自己的数据集

YOLOv5加入注意力机制可分为以下三个步骤：
1.common.py中加入注意力模块
2.yolo.py中增加判断条件
3.yaml文件中添加相应模块

一、CBAM注意力机制添加

（1）在common.py中添加可调用的CBAM模块

1.打开models文件夹中的common.py文件
在这里插入图片描述

2.将下面的CBAMC3代码复制粘贴到common.py文件中

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.sharedMLP = nn.Sequential(
        # nn.Conv2d(in_planes, in_planes // ratio, 1, bias=False), nn.ReLU(),
        # nn.Conv2d(in_planes // rotio, 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 torch.mul(x, 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.conv = 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

最低0.47元/天解锁文章