YOLOv8模型改进2【增加CBAM注意力】

原创

已于 2024-03-13 20:03:37 修改 · 1.2k 阅读

12 ·

CC 4.0 BY-SA版权

文章标签：

#YOLO #人工智能 #python #yolov8 #深度学习

于 2024-03-07 20:55:27 首次发布

本文介绍了如何将CBAM注意力机制应用到YOLOv8模型中，以提升目标检测性能。CBAM结合空间和通道注意力，通过在YOLOv8的关键位置集成，实现在验证集和测试集上的良好表现。文章提供了代码修改路径和训练过程的说明。

一、CBAM注意力机制简介

CBAM注意力机制（Convolutional Block Attention Module）。它结合了空间注意力机制（Spatial Attention Module）以及通道注意力机制(Channel Attention Module)，以此来提升图片分类以及目标检测的结果。

它与SEnet的区别是，SEnet注意力机制只关注通道信息，CBAM注意力机制在通道信息的基础上还增加了对于空间信息的处理。

CBAM结构图如下图所示：（输入特征经过通道注意力机制模块之后，再经过空间注意力机制模块，得到输出特征）
在这里插入图片描述
通道注意力机制模块的结构图如下图所示：

空间注意力机制模块的结构图如下图所示：
在这里插入图片描述

将CBAM模块集成到ResBlock中的结构图如下所示：
在这里插入图片描述

目标检测任务中，CBAM模块在MS COCO验证集上的表现如下：
在这里插入图片描述
目标检测任务中，CBAM模块在VOC 2007测试集上的表现如下：

【注意：1】具体的细节请移步论文：https://arxiv.org/pdf/1807.06521.pdf
【注意：2】GIthub代码:https://github.com/Jongchan/attention-module/blob/master/MODELS/cbam.py

二、增加CBAM注意力机制到YOLOv8模型上

(写在前面：增加CBAM注意力机制的方法与增加SE的方法是一样的，区别在于，YOLOv8原来就自带CBAM模块，所以修改的地方会少一些)

**修改的还是下面的四个地方。

【1: …/ultralytics/nn/modules/conv.py】

conv.py中该增加CBAM模块的地方都已经增加了，可以不用管，或者你检查一下，如果缺少CBAM的代码，请加入以下内容，然后在__all__ =中声明一下，具体的可以对照一下这篇文章，加深一下印象：**https://blog.youkuaiyun.com/A__MP/article/details/136515383?spm=1001.2014.3001.5502

class CBAM(nn.Module):
    """Convolutional Block Attention Module."""

    def __init__(self, c1, kernel_size=7):
        """Initialize CBAM with given input channel (c1) and kernel size."""
        super().__init__()
        self.channel_attention = ChannelAttention(c1)
        self.spatial_attention = SpatialAttention(kernel_size)

    def forward(self, x):
        """Applies the forward pass through C1 module."""
        return self.spatial_attention(self.channel_attention(x))


class Concat(nn.Module):
    """Concatenate a list of tensors along dimension."""

    def __init__(self, dimension=1):
        """Concatenates a list of tensors along a specified dimension."""
        super().__init__()
        self.d = dimension

    def forward(self, x):
        """Forward pass for the YOLOv8 mask Proto module."""
        return torch.cat(x, self.d)