YOLO引入注意力模块

这里对自己简单的修改YOLO网络做一个记录，YOLO这个框架非常成熟（修改yolov8、v9、v11等等都是相同的方法），进行这种即插即用模块的修改非常简单，下面我会以非常典型的CBAM模块为例子。

一、注意力模块

CBAM模块网上介绍的博客非常多，这里我不再赘述。源码如下，要明白的是这个模块不会改变通道数，可以理解为只矫正所学特征。所以才被归为即插即用的模块。另外需要注意这里有个输入通道的参数需要指定，要根据插入到YOLO网络位置修改。

import torch
import torch.nn as nn
 
 
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))

二、添加方式

2.1 模块源码实现

在ultralytics/nn文件路径下，新建attention文件夹，用来存放我们需要添加的所有注意力模块的源码。然后新建一个CBAM.py，把上面CBAM的源码实现copy进去。还需要新建一个\init.py来解决导包时相对路径等等的问题

2.2 将模块导入到task文件中

在ultralytics/nn文件夹下的task.py中，首先导入attention文件夹下所有的模块

直接在task.py中Ctrl+F搜索parse，找到parse_model这个函数，在elif m is AIFI这条语句之前，添加要导入的模块，在这个字典里面把你要插入的模块都添加进去，用逗号分割，比如{CBAM, SE, GAM}，我这里作为演示就只添加了CBAM（注释配合后面的修改网络yaml才能看懂，回过头来再看）。

2.3 修改YOLOv11网络配置文件

YOLO这个框架比较成熟，修改YOLOv8、v9、v10，v11等等都是就修改对应的配置文件即可，这里以YOLOv11为例，我这里是做目标检测，所以直接复制ultralytics/cfg/models/11中的yolo11.yaml文件，重命名为yolo11n-CBAM.yaml。命名需要注意模型的型号n、s、m不要漏，我们都知道YOLO有n、s、m、l，x五个大小型号的网络，而我们修改网络结构后，再通过yaml文件加载网络结构时，模型大小会根据文件名中的型号来确定，我这里是在yolo11n最小的模型上插入CBAM进行实验，所以这里命名为yolo11n-CBAM.yaml。

接下来就是正式修改网络配置文件，这里有两种常见的插入注意力的方法，一是插入到backbone特征提取结束的SPPF层后，另一种是插入到Detect检测头前进行再次矫正。

怎么修改这个网络呢？首先要读懂每一列的意思，第一列表示从哪里连接过来（如果是-1表示上一层)，第三列是模块名称，第二列是这个模块的个数，第四列是这个模块的参数列表，所以如果我们在某一层嵌入一个模块后，就必须修改后续层的第一列的层号，否则就会网络的连接就乱套了，以第一种插入到SPPF后为例，CBAM模块是插入到第10层，from这一列往下检查原本在第10层及其后的层号都需要加1，比如某层本来是concat其上一层和13层做特征融合，就要修改为连接其上一层和14层，以此类推，原来修改后的yaml文件如下：

# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license

# Ultralytics YOLO11 object detection model with P3/8 - P5/32 outputs
# Model docs: https://docs.ultralytics.com/models/yolo11
# Task docs: https://docs.ultralytics.com/tasks/detect

# Parameters
nc: 80 # number of classes
scales: # model compound scaling constants, i.e. 'model=yolo11n.yaml' will call yolo11.yaml with scale 'n'
  # [depth, width, max_channels]
  n: [0.50, 0.25, 1024] # summary: 181 layers, 2624080 parameters, 2624064 gradients, 6.6 GFLOPs
  s: [0.50, 0.50, 1024] # summary: 181 layers, 9458752 parameters, 9458736 gradients, 21.7 GFLOPs
  m: [0.50, 1.00, 512] # summary: 231 layers, 20114688 parameters, 20114672 gradients, 68.5 GFLOPs
  l: [1.00, 1.00, 512] # summary: 357 layers, 25372160 parameters, 25372144 gradients, 87.6 GFLOPs
  x: [1.00, 1.50, 512] # summary: 357 layers, 56966176 parameters, 56966160 gradients, 196.0 GFLOPs

# YOLO11n backbone
backbone:
  # [from, repeats, module, args]
  - [-1, 1, Conv, [64, 3, 2]] # 0-P1/2，64表示conv的输出通道数，3是卷积核大小，2是步长
  - [-1, 1, Conv, [128, 3, 2]] # 1-P2/4
  - [-1, 2, C3k2, [256, False, 0.25]]
  - [-1, 1, Conv, [256, 3, 2]] # 3-P3/8
  - [-1, 2, C3k2, [512, False, 0.25]]
  - [-1, 1, Conv, [512, 3, 2]] # 5-P4/16
  - [-1, 2, C3k2, [512, True]]
  - [-1, 1, Conv, [1024, 3, 2]] # 7-P5/32
  - [-1, 2, C3k2, [1024, True]]
  - [-1, 1, SPPF, [1024, 5]] # 9
  - [-1, 1, CBAM, []] # CBAM模块只有一个输入通道参数，这直接省略会默认为上一层的输出通道数
  - [-1, 2, C2PSA, [1024]] # 10

# YOLO11n head
head:
  - [-1, 1, nn.Upsample, [None, 2, "nearest"]]
  - [[-1, 6], 1, Concat, [1]] # cat backbone P4
  - [-1, 2, C3k2, [512, False]] # 13

  - [-1, 1, nn.Upsample, [None, 2, "nearest"]]
  - [[-1, 4], 1, Concat, [1]] # cat backbone P3
  - [-1, 2, C3k2, [256, False]] # 16 (P3/8-small)

  - [-1, 1, Conv, [256, 3, 2]]
  - [[-1, 14], 1, Concat, [1]] # cat head P4
  - [-1, 2, C3k2, [512, False]] # 19 (P4/16-medium)

  - [-1, 1, Conv, [512, 3, 2]]
  - [[-1, 11], 1, Concat, [1]] # cat head P5
  - [-1, 2, C3k2, [1024, True]] # 22 (P5/32-large)

  - [[17, 20, 23], 1, Detect, [nc]] # Detect(P3, P4, P5)

另一种是插入到Detect检测头前进行再次矫正，就留给大家自己思考和动手操作了，只要明白了上面的例子，就可以开始各种魔改网络了。

2.4 加载修改后的网络并训练（可选加载预训练权重）

主要是想说一下，即使修改后也可以加载原先的预训练权重，之前更新到yolov8的时候要实现加载预训练权重还是有点小麻烦，现在框架可以直接load预训练权重，不适配的层会直接跳过。

import os
from ultralytics import YOLO

model = YOLO("/workspace/ultralytics/ultralytics/cfg/models/11/yolo11n-CBAM.yaml")
model.load("yolo11n.pt") # 加载预训练权重
    # 开始训练
results = model.train(
    data="data/rubbish.yaml",
    epochs=500,
    batch=32,
    imgsz=640,
    name="train_yolov11n_CBAM_loc2_Pre",  # 实验名称
    save=True,  # 保存训练结果
    plots=True,  
    device="0",  
    workers=4,  # 数据加载工作进程数
    patience=100,  # 早停
)