YOLOv10改进-----配置SE注意力机制

什么是SE注意力机制？

在深度学习领域，SE（Squeeze-and-Excitation）模块是一种引人注目的注意力机制，它以其在自适应特征提取中的出色表现而备受关注。SE模块最初在2017年由Jie Hu等人提出，引用量高达2万，旨在改善卷积神经网络（CNN）中通道间信息的传递效率。

在传统的CNN架构中，通常采用卷积层和池化层来提取图像特征。然而，这种方法并没有明确地对特征通道之间的关系进行建模，导致一些通道对于特定任务的贡献相对较小，而其他通道则更重要。SE模块旨在解决这个问题。

SE模块通过引入一个Squeeze操作和一个Excitation操作来建模通道之间的关系。在Squeeze阶段，它通过全局平均池化操作将卷积层的输出特征图压缩成一个特征向量。然后，在Excitation阶段，通过使用全连接层和非线性激活函数，学习生成一个通道的权重向量。这个权重向量被应用于原始特征图上的每个通道，以对不同通道的特征进行加权。

通过这种方式，SE模块能够自适应地学习到每个通道的重要性，并且根据任务的需要加权调整特征图中的通道贡献。这种注意力机制有助于网络更好地关注重要的特征通道，从而提高模型性能。

总而言之，SE模块通过引入Squeeze和Excitation操作，通过自适应地学习每个通道的权重，增强了神经网络的表达能力和性能，使网络能够更有效地学习和利用特征通道之间的关系。它能够自动学习每个通道的权重和重要性。通过学习到的权重，SE模块可以自适应地调整每个通道的特征在后续网络层中的重要性。这样，网络可以更加关注重要的特征通道，提高特征的判别能力，从而提升模型的性能。

在YOLOv10中添加注意力机制有什么好处?

将(SE)机制添加到YOLOv10可以提供以下几个好处:

1. 增强的特征表示:SE通过重新加权通道重要性来帮助模型关注最相关的特征，改善整体特征表示。

2. 提高精度*:通过强调重要通道，模型可以获得更好的检测性能，特别是在有遮挡或杂波的复杂场景中。

3. 自适应学习:SE允许网络自适应地学习哪些特征对特定任务是重要的，从而产生更稳健的预测。

4. 更好的泛化:改进的注意力机制可以帮助模型更好地泛化不同的数据集和场景。

5. 减少过拟合:通过关注基本特征，SE可以帮助减轻过拟合，特别是在训练数据有限的情况下。

给YOLOv10n添加SEAttention

SEAttention源码

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



class SEAttention(nn.Module):

    def __init__(self, channel=512,reduction=16):
        super().__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.fc = nn.Sequential(
            nn.Linear(channel, channel // reduction, bias=False),
            nn.ReLU(inplace=True),
            nn.Linear(channel // reduction, channel, bias=False),
            nn.Sigmoid()
        )


    def init_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                init.kaiming_normal_(m.weight, mode='fan_out')
                if m.bias is not None:
                    init.constant_(m.bias, 0)
            elif isinstance(m, nn.BatchNorm2d):
                init.constant_(m.weight, 1)
                init.constant_(m.bias, 0)
            elif isinstance(m, nn.Linear):
                init.normal_(m.weight, std=0.001)
                if m.bias is not None:
                    init.constant_(m.bias, 0)

    def forward(self, x):
        b, c, _, _ = x.size()
        y = self.avg_pool(x).view(b, c)
        y = self.fc(y).view(b, c, 1, 1)
        return x * y.expand_as(x)

在YOLOv10n主干中添加SEAttention

在ultarlytics/nn/ 路径下新建SEAttention.py文件 将源代码复制进去

修改tultarlytics/nn/ 目录下tasks.py文件，将SEAttention引入主干

在文件的开头引入SEAttention

修改  parse_model函数，找到文件中827行的parse_model函数，在函数下进行修改

        elif m in {SEAttention}:
            args = [ch[f], *args]

修改yolov10n.yaml文件

nc：数据集的分类个数，大家根据自己的数据集进行修改。

这里将SEAttention添加在cf2模块后面，SEAttention的位置不固定，大家根据需求进行添加。

#仅供学习交流
# Parameters
nc: 9 # 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] 

# 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, SCDown, [512, 3, 2]] # 5-P4/16
  - [-1, 6, C2f, [512, True]]
  - [-1, 1, SCDown, [1024, 3, 2]] # 7-P5/32
  - [-1, 3, C2f, [1024, True]]
  - [-1, 1, SEAttention, [1024]] # 9
  - [-1, 1, SPPF, [1024, 5]] # 10
  - [-1, 1, PSA, [1024]] # 11

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

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

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

  - [-1, 1, SCDown, [512, 3, 2]]
  - [[-1, 11], 1, Concat, [1]] # cat head P5
  - [-1, 3, C2fCIB, [1024, True, True]] # 23 (P5/32-large)

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

 在主文件夹目录下新建yaml文件，复制代码

 训练YOLOv10

一定要使用train.py进行训练不要用yolo指令会报错，train.py代码如下：

from ultralytics import YOLOv10

model_yaml_path="yolov10n-SE.yaml"
# 使用你自己的数据集路径
data_yaml_path="D:\\xgk_project\\yolov10\\code\yolov10-main\\yolov10-main\\data\\data.yaml" 

if __name__=="__main__":
    model=YOLOv10(model_yaml_path)
    results=model.train(data=data_yaml_path,
                        epochs=500,
                        batch=16,
                        name="exp")

结果展示

运行train.py，会打印出YOLOv10n的网络结构。

 可以看见SE已经被添加在网络中啦~，希望能帮助到你。