YOLO11改进 | 自定义数据集训练 | VMamba助力YOLO11检测（附完整代码 + 修改教程 + 运行教程+个人数据标注教程+半自动标注教程+训练结果可视化教程+部署教程）

机＿长

已于 2024-10-08 14:28:59 修改

阅读量368

点赞数 21

文章标签： yolo11 VMamba YOLO

于 2024-10-08 14:28:02 首次发布

本文链接：https://blog.csdn.net/weixin_51988935/article/details/142757964

版权

一、本文介绍

二、vision mamba结构与原理

三、yolov11与VMamba结合修改教程

3.1 核心代码文件的创建与添加

3.1.1 vmamba.py文件添加

3.1.2 __init__.py文件添加

3.2 yolov11 的tasks.py 文件修改

3.2.1 导入vmamba模块

3.2.2 修改并添加vmamba模块

3.3 模型配置文件添加

3.3.1 构建模型配置文件yolo11-vmambanet.yaml

6.1.加载训练权重进行待标注数据集的推理

6.2.将推理结果txt转换为xml标签文件

6.3.打开标注软件labelimg，微调标注

一、本文介绍

本文提出了一种名为Mamba-UNet的新型网络架构，该架构通过结合U-Net与Mamba架构的优势，旨在提升医学图像分割的性能。在MRI心脏分割数据集上的实验结果表明，Mamba-UNet相较于传统的UNet和先进的SwinUNet，展现出了更优越的性能。此外，本文还详细探讨了如何将Mamba注意力模块应用于Yolo11，创造出Yolo11-Mamba网络，该网络在全局特征捕捉能力上得到了显著提升。此外、还提供将其与YOLO11等目标检测模型有效结合改进的方法，并对自定义数据集的训练以及半自动化标注的方法进行简单介绍，帮助力读者快速在自己的数据集上进行实验与应用。

二、vision mamba结构与原理

Mamba成功的关键在于S6模型，该模型为NLP任务设计，通过选择性扫描空间状态序列模型，将二次复杂度降低至线性。但由于视觉信号（如图像）的无序性，Mamba的S6模型不能直接应用，

设计了 2D-selective-scan（SS2D）模块。

如上图所示，SS2D由三个部分组成：扫描expanding操作、S6块操作和扫描merging操作。如图2(a)所示，扫描expanding操作沿着四个不同的方向(左上到右下、左下到右上、右下到左上、右上到左下)将输入图像展开成序列。然后通过S6块对这些序列进行特征提取，确保各个方向的信息被彻底扫描，从而捕获不同的特征。随后，如图2(b)所示，扫描merging操作将来自四个方向的序列相加并合并，将输出图像恢复为与输入相同的大小。源自Mamba[16]的S6块在S4[17]之上引入了一种选择机制，通过根据输入调整SSM的参数。这使模型能够区分并保留相关信息，同时过滤掉不相关的信息。

上图为本文提出的 VMamba 结构图。VMamba 的整体框架与主流的视觉模型类似，如上图 (b)所示。经过Layer Normalization后，输入被分成两个分支。在第一个分支中，输入经过一个线性层，然后是一个激活函数。在第二个分支中，输入通过线性层、深度可分离卷积和激活函数进行处理，然后输入到2D选择性扫描(SS2D)模块中进行进一步的特征提取。随后，使用Layer Normalization对特征进行归一化，然后使用第一个分支的输出执行逐元素的生成，以合并两条路径。最后，使用线性层混合特征，并将此结果与残差连接相结合，形成VSS块的输出。本文默认采用SiLU作为激活函数。

三、yolov11与VMamba结合修改教程

基础环境的安装

conda create -n vmunet python=3.8
conda activate vmunet
pip install torch==1.13.0 torchvision==0.14.0 torchaudio==0.13.0 --extra-index-url https://download.pytorch.org/whl/cu117
pip install packaging
pip install timm==0.4.12
pip install pytest chardet yacs termcolor
pip install submitit tensorboardX
pip install triton==2.0.0
pip install causal_conv1d==1.0.0  # causal_conv1d-1.0.0+cu118torch1.13cxx11abiFALSE-cp38-cp38-linux_x86_64.whl
pip install mamba_ssm==1.0.1  # mmamba_ssm-1.0.1+cu118torch1.13cxx11abiFALSE-cp38-cp38-linux_x86_64.whl
pip install scikit-learn matplotlib thop h5py SimpleITK scikit-image medpy yacs

3.1 核心代码文件的创建与添加

3.1.1 vmamba.py文件添加

找到ultralytics/nn文件夹，建立Addmoudules文件夹，并在该文件夹下建立vmamba.py文件，如下图。

vmamba.py添加核心代码，如下

class VSSM(nn.Module):
    def __init__(self, patch_size=4, in_chans=3, num_classes=1000, depths=[2, 2, 9, 2], depths_decoder=[2, 9, 2, 2],
                 dims=[96, 192, 384, 768], dims_decoder=[768, 384, 192, 96], d_state=16, drop_rate=0.,
                 attn_drop_rate=0., drop_path_rate=0.1,
                 norm_layer=nn.LayerNorm, patch_norm=True,
                 use_checkpoint=False, **kwargs):
        super().__init__()
        self.num_classes = num_classes
        self.num_layers = len(depths)
        if isinstance(dims, int):
            dims = [int(dims * 2 ** i_layer) for i_layer in range(self.num_layers)]
        self.embed_dim = dims[0]
        self.num_features = dims[-1]
        self.dims = dims

        self.patch_embed = PatchEmbed2D(patch_size=patch_size, in_chans=in_chans, embed_dim=self.embed_dim,
                                        norm_layer=norm_layer if patch_norm else None)

        # WASTED absolute position embedding ======================
        self.ape = False
        # self.ape = False
        # drop_rate = 0.0
        if self.ape:
 
.........

注：全部代码私信博主获取

3.1.2 init.py文件添加

在Addmoudules文件夹建立__init__.py文件，如下图

__init__.py文件添加代码，如下

from .vmamba import *

3.2 yolov11 的tasks.py 文件修改

找到ultralytics/nn文件夹下的task.py 文件并打开

3.2.1 导入vmamba模块

3.2.2 修改并添加vmamba模块

task.py 343 行修改

task.py，1002 添加MambaLayer

注：其中红色框为添加部分

3.3 模型配置文件添加

3.3.1 构建模型配置文件yolo11-vmambanet.yaml

找到ultralytics/cfg/models/v8文件夹，并在该文件夹下建立yolo11-vmambanet.yaml文件，如下图

3.3.2 配置网络结构

# Ultralytics YOLO 🚀, AGPL-3.0 license
# YOLO11 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=yolo11n.yaml' will call yolo11.yaml with scale 'n'
  # [depth, width, max_channels]
  n: [0.50, 0.25, 1024] # summary: 319 layers, 2624080 parameters, 2624064 gradients, 6.6 GFLOPs
  s: [0.50, 0.50, 1024] # summary: 319 layers, 9458752 parameters, 9458736 gradients, 21.7 GFLOPs
  m: [0.50, 1.00, 512] # summary: 409 layers, 20114688 parameters, 20114672 gradients, 68.5 GFLOPs
  l: [1.00, 1.00, 512] # summary: 631 layers, 25372160 parameters, 25372144 gradients, 87.6 GFLOPs
  x: [1.00, 1.50, 512] # summary: 631 layers, 56966176 parameters, 56966160 gradients, 196.0 GFLOPs

# YOLO11n backbone
backbone:
  # [from, repeats, module, args]
  - [ -1, 1, Conv, [ 64, 3, 2 ] ]  # 0-P1/2      # 0.  320
  - [ -1, 1, Conv, [ 128, 3, 2 ] ]  # 1-P2/4     # 1.  160
  - [ -1, 3, MambaLayer, [ 128 ] ]                # 2.  160
  - [ -1, 1, Conv, [ 256, 3, 2 ] ]  # 3-P3/8     # 3.  80
  - [ -1, 6, MambaLayer, [ 256 ] ]                # 4.  80
  - [ -1, 1, Conv, [ 512, 3, 2 ] ]  # 5-P4/16    # 5.  40
  - [ -1, 6, MambaLayer, [ 512 ] ]                # 6.  40
  - [ -1, 1, Conv, [ 1024, 3, 2 ] ]  # 7-P5/32   # 7.  20S
  - [ -1, 3, MambaLayer, [ 1024 ] ]               # 8.  20
  - [ -1, 1, SPPF, [ 1024, 5 ] ]  # 9
  - [-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, 13], 1, Concat, [1]] # cat head P4
  - [-1, 2, C3k2, [512, False]] # 19 (P4/16-medium)

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

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

注：全部代码私信博主获取

四、自定义数据集

4.1数据标注

新建两个文件夹如下，其中images用于保存待标注的图片数据，Annotations用于保存标注过后xml格式的标签文件

打开 labelimg 标注工具（标注工具自行检索下载），如下图，配置images，Annotations文件夹所在路径，然后进行标注

标注后将在Annotations文件夹下生成如下文件

4.2 数据格式转换

如下图新建label文件夹，用于保存由Annotations文件夹下的xml标签文件转换后，适合yolov8训练的txt格式的标签文件

定位至Yolo11_Dataset_Processing/xmls2yolo.py文件并打开，再定位如下图到红色框处，需配置检测的标签列表，以及xml标签路径（Annotations），转换后txt格式的标签路径（label）。

配置完成后运行xmls2yolo.py，即可在label文件夹下生成所需txt格式的标签文件，如下图

注：全部代码联系博主获得

4.3 数据集划分

定位至Yolo11_Dataset_Processing/dataset_split.py文件，打开并定位至红色框处，需配置

classes、source_file_path 、images_file_path ，如下图所示

配置完成后运行dataset_split.py，即可在lmages同级文件夹下生成划分后的训练集、验证集、测试集，如下图

4.4 训练数据集配置文件添加

定位至ultralytics/cfg/datasets文件夹下，新建训练数据集的配置文件a_sleep.yaml如下图

在a_sleep.yaml 文件中进行训练、验证、测试数据的配置，如下图

注：若有多个训练、验证、测试数据，只需在对应位置继续添加即可，参考注释掉的部分

五、模型训练

5.1 训练文件的创建与配置

在ultralytics的同级目录处创建train.py 文件

点击打开train.py文件，如下图需配置 1，与 2处，其中1配置的是本文3.3构建模型配置文件yolov8-vmambanet.yaml。2处配置的是本文4.4 新建训练数据集的配置文件a_sleep.yaml。

5.2 训练过程展示

六、半自动标注

半自动化标注，解决的是标注数据量庞大，难以人工从头进行一一标注的问题。主要步骤是，先标注两千张数据，进行模型训练，然后用训练好的模型，对剩下未标注的数据集进行推理，得到模型推理后的结果，将这些结果转换为xml格式的标签文件。最好利用标注软件labelimg打开，就得到大量已经有预标注的图片，后续工作仅需微调这些标注即可，大大减少了标注的工作量。以下是半自动化标注流程