YOLOv10+ECA，魔改基础

最新推荐文章于 2024-10-12 14:01:20 发布

一口小西瓜

最新推荐文章于 2024-10-12 14:01:20 发布

阅读量799

点赞数 9

分类专栏： YOLOv10魔改基础文章标签： YOLO 深度学习 pytorch python

本文链接：https://blog.csdn.net/m0_73730717/article/details/141301386

版权

YOLOv10魔改基础专栏收录该内容

3 篇文章 0 订阅

订阅专栏

论文地址：1910.03151 (arxiv.org)

代码地址:https://github.com/BangguWu/ECANet

1 ECA简介

总结优点：

ECA事实上是SENET的改进版，它去除了原来SENET中的全连接层，换成了1*1的卷积核进行处理，使得模型参数变小，变得更加轻量级，因为ECA的作者认为卷积具有良好的跨通道信息捕捉能力，因此捕捉所有通道的信息是没必要的，因此取消了全连接层，换成了1*1的卷积

1. 计算高效
由于ECA模块避免了复杂的降维和升维过程，并且使用了简单的一维卷积操作这使得ECA模块能够在不增加显著计算负担的情况下，为模型带来性能提升，提高运算效率。

2. 保留信息完整性
与传统的注意力机制相比，ECA模块无需进行降维和升维的操作，进而保留原始通道特征的信息完整性。这有助于模型更好地利用通道间的依赖关系，提升特征表示能力。

3. 自适应核大小
ECA模块能够根据通道数自适应地调整一维卷积的核大小，使其能够灵活地捕捉不同范围内的通道依赖性。这种自适应机制使得ECA模块在不同规模的网络和不同深度的层次中都能有效工作。

2 YOLOv10部署

2.1 在ultralytics/nn/modules文件下新建ECA.py

将下面代码复制过去：

下面注释的是一些简单调用，确定代码可以正常运行！很简单，正常使用的时候注释就好（下面已经注释）。

"""Constructs a ECA module.
一口小西瓜
"""
import torch.nn as nn
class ECA(nn.Module):
    """Constructs a ECA module.
    Args:
        channel: Number of channels of the input feature map
        k_size: Adaptive selection of kernel size
    """

    def __init__(self, c1, c2, k_size=3):
        super(ECA, self).__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.conv = nn.Conv1d(1, 1, kernel_size=k_size, padding=(k_size - 1) // 2, bias=False)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        # feature descriptor on the global spatial information
        y = self.avg_pool(x)

        # print(y.shape,y.squeeze(-1).shape,y.squeeze(-1).transpose(-1, -2).shape)
        # Two different branches of ECA module
        # 50*C*1*1
        # 50*C*1
        # 50*1*C
        y = self.conv(y.squeeze(-1).transpose(-1, -2)).transpose(-1, -2).unsqueeze(-1)

        # Multi-scale information fusion
        y = self.sigmoid(y)

        return x * y.expand_as(x)


# # 主函数用于测试 ECA 类
# def main():
#     # 设置随机种子以保证结果可复现
#     torch.manual_seed(0)
#
#     # 创建输入数据：假设批大小为 2，通道数为 16，图像大小为 32x32
#     batch_size = 2
#     channels = 16
#     height = 32
#     width = 32
#     input_tensor = torch.randn(batch_size, channels, height, width)  # 随机生成输入张量
#
#     # 创建 EMA 实例
#     eca = ECA(256,256)
#
#     # 前向传播，通过 EMA 模块
#     output_tensor =eca(input_tensor)
#     # 输出结果的形状
#     print("Input shape: ", input_tensor.shape)
#     print("Output shape: ", output_tensor.shape)
#
#
# if __name__ == "__main__":
#     main()

2.2 调用EMA

步骤1：在ultralytics/nn/modules/__init__.py 中添加，添加方式如图

步骤2：在ultralytics/nn/tasks.py 中添加，添加方式如图

在 tasks.py的头上方就是哈！！

步骤3：在tasks.py中ctrl+F查找def parse_model，在图中位置添加即可！

就是在这块：

        n = n_ = max(round(n * depth), 1) if n > 1 else n  # depth gain
        if m in {

到这里就OK啦！！！剩下的就是yaml文件啦

2.3 yaml文件

将ultralytics/cfg/models/v10/yolov10s.yaml复制，重命名就好了，我是叫ECA.yaml

我就直接丢在head中啦！你们想丢哪里都可以，对应添加就好了！

# 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]
  s: [0.33, 0.50, 1024]

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, C2fCIB, [1024, True, True]]
  - [-1, 1, SPPF, [1024, 5]] # 9
  - [-1, 1, PSA, [1024]] # 10

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

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

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

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

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

2.4 训练调用啦

在ultralytics目录下建立train.py,代码如下：按照对应填写即可哦

# coding:utf-8
from ultralytics import YOLOv10

# 模型配置文件
model_yaml_path = "ultralytics/cfg/models/v10/yolov10s-ECA.yaml"
# 数据集配置文件
data_yaml_path = ''
# 预训练模型
pre_model_name = r'权重地址'

if __name__ == '__main__':
    # 加载预训练模型
    model = YOLOv10("ultralytics/cfg/models/v10/yolov10s-ECA.yaml").load(r'权重地址')
    # 训练模型
    results = model.train(data=data_yaml_path, epochs=100, batch=8,  project='runs/train',name='exp')

呐呐呐，打印出来了，可以正常训练！因为数据集不同，结果也不同，还需要尝试！