免费阅读篇 | 芒果YOLOv8改进111：注意力机制CBAM：轻量级卷积块注意力模块，无缝集成到任何CNN架构中，开销可以忽略不计

💡🚀🚀🚀本博客改进源代码改进适用于 YOLOv8 按步骤操作运行改进后的代码即可

该专栏完整目录链接：芒果YOLOv8深度改进教程

该篇博客为免费阅读内容，YOLOv8+CBAM改进内容🚀🚀🚀

文章目录

- - 1. CBAM 论文
  - 2. YOLOv8 核心代码改进部分
  - 2.1 核心新增代码
  - - 2.2 代码修改部分
  - 2.3 YOLOv8-CBAM 网络配置文件
  - 2.4 运行代码
  - 改进说明

1. CBAM 论文

在这里插入图片描述

我们提出了卷积块注意力模块（CBAM），这是一种用于前馈卷积神经网络的简单而有效的注意力模块。给定中间特征图，我们的模块沿着两个独立的维度（通道和空间）顺序推断注意力图，然后将注意力图乘以输入特征图以进行自适应特征细化。由于 CBAM 是一个轻量级通用模块，因此它可以无缝集成到任何 CNN 架构中，且开销可以忽略不计，并且可以与基础 CNN 一起进行端到端训练。我们通过在 ImageNet-1K、MS~COCO 检测和 VOC~2007 检测数据集上进行大量实验来验证我们的 CBAM。我们的实验表明各种模型的分类和检测性能得到了一致的改进，证明了 CBAM 的广泛适用性。代码和模型将公开。

在这里插入图片描述

具体细节可以去看原论文：https://arxiv.org/pdf/1807.06521.pdf

2. YOLOv8 核心代码改进部分

2.1 核心新增代码

首先在ultralytics/nn/modules文件夹下，创建一个 cbam.py文件，新增以下代码

import numpy as np
import torch
from torch import nn
from torch.nn import init

class ChannelAttentionModule(nn.Module):
    def __init__(self, c1, reduction=16):
        super(ChannelAttentionModule, self).__init__()
        mid_channel = c1 // reduction
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.max_pool = nn.AdaptiveMaxPool2d(1)

        self.shared_MLP = nn.Sequential(
            nn.Linear(in_features=c1, out_features=mid_channel),
            nn.LeakyReLU(0.1, inplace=True),
            nn.Linear(in_features=mid_channel, out_features=c1)
        )
        self.act = nn.Sigmoid()
        #self.act=nn.SiLU()
    def forward(self, x):
        avgout = self.shared_MLP(self.avg_pool(x).view(x.size(0),-1)).unsqueeze(2).unsqueeze(3)
        maxout = self.shared_MLP(self.max_pool(x).view(x.size(0),-1)).unsqueeze(2).unsqueeze(3)
        return self.act(avgout + maxout)

class SpatialAttentionModule(nn.Module):
    def __init__(self):
        super(SpatialAttentionModule, self).__init__()
        self.conv2d = nn.Conv2d(in_channels=2, out_channels=1, kernel_size=7, stride=1, padding=3)
        self.act = nn.Sigmoid()
    def forward(self, x):
        avgout = torch.mean(x, dim=1, keepdim=True)
        maxout, _ = torch.max(x, dim=1, keepdim=True)
        out = torch.cat([avgout, maxout], dim=1)
        out = self.act(self.conv2d(out))
        return out

class CBAM(nn.Module):
    def __init__(self, c1,c2):
        super(CBAM, self).__init__()
        self.channel_attention = ChannelAttentionModule(c1)
        self.spatial_attention = SpatialAttentionModule()

    def forward(self, x):
        out = self.channel_attention(x) * x
        out = self.spatial_attention(out) * out
        return out

2.2 代码修改部分

在ultralytics/nn/modules/init.py中导入定义在 cbam.py 里面的模块

from .cbam import CBAM

'CBAM' 加到 __all__ = [...] 里面

第一步：
在ultralytics/nn/tasks.py文件中,新增

from ultralytics.nn.modules import CBAM

直接复制这段代码即可

如下图所示：

在这里插入图片描述

然后在在tasks.py中配置
找到

        elif m is nn.BatchNorm2d:
        args = [ch[f]]

在这句上面加一个

        elif m is CBAM:
            c1, c2 = ch[f], args[0]
            if c2 != nc:  # if c2 not equal to number of classes (i.e. for Classify() output)
                c2 = make_divisible(min(c2, max_channels) * width, 8)

            args = [c1, c2, *args[1:]]

直接复制这段代码即可

如下图所示
在这里插入图片描述

2.3 YOLOv8-CBAM 网络配置文件

新增YOLOv8-CBAM.yaml

# Ultralytics YOLO 🚀, AGPL-3.0 license
# YOLOv8 object detection model with P3-P5 outputs. For Usage examples see https://docs.ultralytics.com/tasks/detect

# Parameters
nc: 80  # number of classes
scales: # model compound scaling constants, i.e. 'model=yolov8n.yaml' will call yolov8.yaml with scale 'n'
  # [depth, width, max_channels]
  n: [0.33, 0.25, 1024]  # YOLOv8n summary: 225 layers,  3157200 parameters,  3157184 gradients,   8.9 GFLOPs
  s: [0.33, 0.50, 1024]  # YOLOv8s summary: 225 layers, 11166560 parameters, 11166544 gradients,  28.8 GFLOPs
  m: [0.67, 0.75, 768]   # YOLOv8m summary: 295 layers, 25902640 parameters, 25902624 gradients,  79.3 GFLOPs
  l: [1.00, 1.00, 512]   # YOLOv8l summary: 365 layers, 43691520 parameters, 43691504 gradients, 165.7 GFLOPs
  x: [1.00, 1.25, 512]   # YOLOv8x summary: 365 layers, 68229648 parameters, 68229632 gradients, 258.5 GFLOPs

# YOLOv8.0n backbone
backbone:
  # [from, repeats, module, args]
  - [-1, 1, Conv, [64, 3, 2]]  # 0-P1/2
  - [-1, 1, Conv, [128, 3, 2]]  # 1-P2/4
  - [-1, 3, C2f, [128, True]]
  - [-1, 1, Conv, [256, 3, 2]]  # 3-P3/8
  - [-1, 6, C2f, [256, True]]
  - [-1, 1, Conv, [512, 3, 2]]  # 5-P4/16
  - [-1, 6, C2f, [512, True]]
  - [-1, 1, Conv, [1024, 3, 2]]  # 7-P5/32
  - [-1, 3, C2f, [1024, True]]
  - [-1, 3, CBAM, [1024]]
  - [-1, 1, SPPF, [1024, 5]]  # 9

# YOLOv8.0n head
head:
  - [-1, 1, nn.Upsample, [None, 2, 'nearest']]
  - [[-1, 6], 1, Concat, [1]]  # cat backbone P4
  - [-1, 3, C2f, [512]]  # 12

  - [-1, 1, nn.Upsample, [None, 2, 'nearest']]
  - [[-1, 4], 1, Concat, [1]]  # cat backbone P3
  - [-1, 3, C2f, [256]]  # 15 (P3/8-small)

  - [-1, 1, Conv, [256, 3, 2]]
  - [[-1, 13], 1, Concat, [1]]  # cat head P4
  - [-1, 3, C2f, [512]]  # 18 (P4/16-medium)

  - [-1, 1, Conv, [512, 3, 2]]
  - [[-1, 10], 1, Concat, [1]]  # cat head P5
  - [-1, 3, C2f, [1024]]  # 21 (P5/32-large)

  - [[16, 19, 22], 1, Detect, [nc]]  # Detect(P3, P4, P5)