YOLO11 改进 – 注意力机制 _ CBAM (Convolutional Block Attention Module) 卷积块注意力模块:轻量级设计自适应优化特征,提升小目标检测精度

前言

本文介绍了卷积块注意力模块(CBAM)及其在YOLOv11中的集成。CBAM是一种用于CNN的注意力机制,通过依次在通道和空间维度推断注意力图,自适应优化输入特征图。该模块轻量通用,可无缝集成到任何CNN架构中。我们将CBAM引入YOLOv11,在检测头部分的不同尺度特征图上应用该模块。实验表明,改进后的YOLOv11在目标检测任务中表现良好,证明了CBAM在提升模型性能方面的有效性和广泛适用性。

文章目录: YOLOv11改进大全:卷积层、轻量化、注意力机制、损失函数、Backbone、SPPF、Neck、检测头全方位优化汇总

专栏链接: YOLOv11改进专栏

文章目录

[TOC]

介绍

摘要

我们提出了一种名为卷积块注意力模块(Convolutional Block Attention Module,CBAM)的新型注意力机制,该模块设计简洁但效果显著,专门用于前馈卷积神经网络。在给定中间特征图的情况下,该模块依次沿着通道和空间两个独立维度推断注意力图,随后将所得注意力图与输入特征图进行乘法操作以实现自适应特征优化。CBAM作为一种轻量级且通用性强的模块,能够无缝集成到任意卷积神经网络(CNN)架构中,其计算开销可忽略不计,并支持与基础CNN进行端到端的联合训练。通过在ImageNet-1K、MS COCO检测以及VOC 2007检测数据集上进行广泛实验,我们验证了CBAM的有效性。实验结果表明,该模块在多种模型架构中均能带来分类和检测性能的一致提升,充分证明了其广泛的适用性。相关代码与模型将予以公开发布。

文章链接

论文地址: 论文地址

代码地址: 代码地址

参考代码: yolov8源代码中的 ultralytics.nn.modules.conv.CBAM

基本原理

Convolutional Block Attention Module(CBAM)是一种针对卷积神经网络(CNN)设计的新型注意力机制,旨在增强其在各种计算机视觉任务中的性能,如图像分类和目标检测。CBAM依次沿着通道和空间两个独立维度推断注意力图,然后将这些图结合起来,自适应地优化输入特征图。

  1. 通道注意力模块 :CBAM中的通道注意力机制专注于建模特征图中不同通道之间的相互关系。它学习为每个通道分配不同的重要性权重,通过捕捉通道间的关系,有助于网络强调信息丰富的通道,同时抑制不太相关的通道,从而改善特征表示。

  2. 空间注意力模块 :CBAM中的空间注意力机制旨在捕捉特征图中每个通道内的空间依赖关系。通过分析不同位置之间的空间关系,空间注意力模块学习突出重要的空间区域,同时抑制不相关的区域。这使网络能够专注于相关的空间位置,以更好地提取特征。

以下为 CBAM (Convolutional Block Attention Module) 的架构

从架构来看,其实蛮像 BAM(Bottleneck Attention Module),一样有 Channel attention 和 Spatial attention 的模块在里面,而 CBAM (Convolutional Block Attention Module) 的流程如下:

  1. 输入特征 通过 Channel Attention 后得到 Mc(F),Mc(F) 与 F(input feature) 做 Element-wise Product 得到 F’

  2. F’ 通过空间注意力 后得到 Ms(F),Ms(F) 与 F’(input feature) 做 Element-wise Product 得到 F’’

最终的输出为 F’’ = Ms(F’) ⊗F’ , F’ = Mc(F) ⊗ F (Input feature) 可以用以下公式来表示

核心代码


class ChannelAttention(nn.Module):
    """Channel-attention module https://github.com/open-mmlab/mmdetection/tree/v3.0.0rc1/configs/rtmdet."""

    def __init__(self, channels: int) -> None:
        """Initializes the class and sets the basic configurations and instance variables required."""
        super().__init__()
        self.pool = nn.AdaptiveAvgPool2d(1)
        self.fc = nn.Conv2d(channels, channels, 1, 1, 0, bias=True)
        self.act = nn.Sigmoid()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Applies forward pass using activation on convolutions of the input, optionally using batch normalization."""
        return x * self.act(self.fc(self.pool(x)))

class SpatialAttention(nn.Module):
    """Spatial-attention module."""

    def __init__(self, kernel_size=7):
        """Initialize Spatial-attention module with kernel size argument."""
        super().__init__()
        assert kernel_size in (3, 7), "kernel size must be 3 or 7"
        padding = 3 if kernel_size == 7 else 1
        self.cv1 = nn.Conv2d(2, 1, kernel_size, padding=padding, bias=False)
        self.act = nn.Sigmoid()

    def forward(self, x):
        """Apply channel and spatial attention on input for feature recalibration."""
        return x * self.act(self.cv1(torch.cat([torch.mean(x, 1, keepdim=True), torch.max(x, 1, keepdim=True)[0]], 1)))

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))

实验

脚本

import warnings
warnings.filterwarnings('ignore')
from ultralytics import YOLO

if __name__ == '__main__':
#     修改为自己的配置文件地址
    model = YOLO('/root/ultralytics-main/ultralytics/cfg/models/11/yolov11-CBAM.yaml')
#     修改为自己的数据集地址
    model.train(data='/root/ultralytics-main/ultralytics/cfg/datasets/coco8.yaml',
                cache=False,
                imgsz=640,
                epochs=10,
                single_cls=False,  # 是否是单类别检测
                batch=8,
                close_mosaic=10,
                workers=0,
                optimizer='SGD',
                amp=True,
                project='runs/train',
                name='CBAM',
                )

结果

THE END