YOLOV8改进:增加注意力模块,以CBAM模块为例

1,conv.py修改

在ultralytics\nn\modules\conv.py该文件中，我们可以写入自己的注意力模块，或者使用V8已经提供的CBAM模块（见代码的CBAM类）,其实yolov8中已经加入了CBAM

第一部在__all__中加入CBAM(模块名)

__all__ = (
    "Conv",
    "Conv2",
    "LightConv",
    "DWConv",
    "DWConvTranspose2d",
    "ConvTranspose",
    "Focus",
    "GhostConv",
    "ChannelAttention",
    "SpatialAttention",
    "CBAM", #🚀
    "Concat",
    "RepConv",
    "BiFPN_Concat2",
    "BiFPN_Concat3",
    "RFACA",
)

第二步；在Conv.py中加入CBAM模块代码

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

2，__init__.py修改

添加完毕后在ultralytics\nn\modules\__init__.py中导入模块，一是在from.conv import中导入。二是在__all__=（）中添加

导包

from .conv import (
    CBAM,
    ChannelAttention,
    Concat,
    Conv,
    Conv2,
    ConvTranspose,
    DWConv,
    DWConvTranspose2d,
    Focus,
    GhostConv,
    LightConv,
    RepConv,
    SpatialAttention,
    BiFPN_Concat2,
    BiFPN_Concat3,
    RFACA,
)

加入

__all__ = (
    "Conv",
    "Conv2",
    "LightConv",
    "RepConv",
    "DWConv",
    "DWConvTranspose2d",
    "ConvTranspose",
    "Focus",
    "GhostConv",
    "ChannelAttention",
    "SpatialAttention",
    "CBAM",
    ）

3.tasks.py修改

第一在from ultralytics.nn.models import中添加模块

from ultralytics.nn.modules import (
    AIFI,
    C1,
    C2,
    C3,
    C3TR,
    OBB,
    SPP,
    SPPF,
    Bottleneck,
    BottleneckCSP,
    C2f,
    C2fAttn,
    CBAM，
）

第二在def parse_model函数中激活---添加分支elif m is CBAM:，具体代码如下：

        elif m in {CBAM}:
            c1, c2 = ch[f], args[0]
            if c2 != nc:
               c2 = make_divisible(min(c2,max_channels) * width,8)
            args = [c1, *args[1:]]

到此就把cbam模块添加完毕

4配置文件xx.yaml

样例如下，放入位置记得对应修改通道数和层数，另外就是CBAM注意力机制的卷积核大小为3/7，其他报错，示例如下

# YOLOAir 🚀 by 🥭, GPL-3.0 license
 
# Parameters      
nc: 11  # number of classes     V5M     3（masks face license plate）
 
scales:
  #n: [0.33, 0.25, 1024]
  s: [0.33, 0.50, 1024]
  m: [0.67, 0.75, 768]
  l: [1.00, 1.00, 512]
  x: [1.00, 1.25, 512]

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, [768, 3, 2]] # 7-P5/32
  - [-1, 3, C2f, [768, True]]
  - [-1, 1, Conv, [1024, 3, 2]] # 9-P6/64
  - [-1, 3, C2f, [1024, True]]

  - [-1,3,CBAM,[1024,7]]  ##11🥭
  - [-1, 1, SPPF, [1024, 5]] # 11
  
# YOLOv8.0x6 head
head:
  - [-1, 1, nn.Upsample, [None, 2, "nearest"]]
  - [[-1, 8], 1, BiFPN_Concat2, [1]] # cat backbone P5
  - [-1, 3, C2f, [768, False]] # 14

  - [-1, 1, nn.Upsample, [None, 2, "nearest"]]
  - [[-1, 6], 1, BiFPN_Concat2, [1]] # cat backbone P4
  - [-1, 3, C2f, [512, False]] # 17

  - [-1, 1, nn.Upsample, [None, 2, "nearest"]]
  - [[-1, 4], 1, BiFPN_Concat2, [1]] # cat backbone P3
  - [-1, 3, C2f, [256, False]] # 20 (P3/8-small)

  - [-1, 1, Conv, [256, 3, 2]]
  - [[-1, 18], 1, BiFPN_Concat2, [1]] # cat head P4    +1
  - [-1, 3, C2f, [512, False]] # 23 (P4/16-medium)

  - [-1, 1, Conv, [512, 3, 2]]
  - [[-1, 15], 1, BiFPN_Concat2, [1]] # cat head P5     +1
  - [-1, 3, C2f, [768, False]] # 26 (P5/32-large)

  - [-1, 1, Conv, [768, 3, 2]]
  - [[-1, 12], 1,BiFPN_Concat2, [1]] # cat head P6        Concat   +1        ，
  - [-1, 3, C2f, [1024, False]] # 29 (P6/64-xlarge)                   c2模块不适合拟合

  - [[21, 24, 27, 30], 1, Detect, [nc]] # Detect(P3, P4, P5, P6)   +1

运行训练打印出完整模型即成功

增加对应的模块后，之后的层数的layer+1，因此需要适当更改，不然会报concat维度不匹配的错误，如下

RuntimeError: Sizes of tensors must match except in dimension 1. Expected size 16 but got size 32 for tensor number 1 in the list.