YOLOv11 数据集训练、推理及网络结构详解

代码需要修改的地方如下：

1. postfix 参数填图片的后缀，需要注意图片格式要统一，是 png 格式就写 png，是 jpg 格式就写 jpg
2. imgpath 参数填图片所在的路径
3. xmlpath 参数填标注文件的路径
4. txtpath 参数填生成的 yolo 格式的文件路径

3. 数据集划分

划分训练集和验证集代码如下：

import os, shutil
from sklearn.model_selection import train_test_split

val_size = 0.2
postfix = 'jpg'
imgpath = r'./data/images'
txtpath = r'./data/labels'
output_train_img_folder = r'./dataset_kengwa/images/train'
output_val_img_folder = r'./dataset_kengwa/images/val'
output_train_txt_folder = r'./dataset_kengwa/labels/train'
output_val_txt_folder = r'./dataset_kengwa/labels/val'

os.makedirs(output_train_img_folder, exist_ok=True)
os.makedirs(output_val_img_folder, exist_ok=True)
os.makedirs(output_train_txt_folder, exist_ok=True)
os.makedirs(output_val_txt_folder, exist_ok=True)

listdir = [i for i in os.listdir(txtpath) if 'txt' in i]
train, val = train_test_split(listdir, test_size=val_size, shuffle=True, random_state=0)

for i in train:
    img_source_path = os.path.join(imgpath, '{}.{}'.format(i[:-4], postfix))
    txt_source_path = os.path.join(txtpath, i)
    img_destination_path = os.path.join(output_train_img_folder, '{}.{}'.format(i[:-4], postfix))
    txt_destination_path = os.path.join(output_train_txt_folder, i)
    shutil.copy(img_source_path, img_destination_path)
    shutil.copy(txt_source_path, txt_destination_path)

for i in val:
    img_source_path = os.path.join(imgpath, '{}.{}'.format(i[:-4], postfix))
    txt_source_path = os.path.join(txtpath, i)
    img_destination_path = os.path.join(output_val_img_folder, '{}.{}'.format(i[:-4], postfix))
    txt_destination_path = os.path.join(output_val_txt_folder, i)
    shutil.copy(img_source_path, img_destination_path)
    shutil.copy(txt_source_path, txt_destination_path)

下面四个参数只需在自己电脑任意位置新建一个文件夹就行，用于存放生成的训练集和验证集。

数据集有以下两种方式放置，都可以进行训练，常见的数据集放置是第一种。

4. 修改 YOLO 的训练配置文件

我们需要在项目下创建一个 data.yaml 的文件，文件名根据数据集名称取。

代码如下：

train: ./images/train # train images (relative to 'path')
val: ./images/val # val images (relative to 'path')
nc: 2 # class names
names: ['dog', 'cat']

三、YOLO 环境配置教程

YOLOv11/YOLOv10/YOLOv9/YOLOv8/YOLOv7/YOLOv5 环境都是通用的，只需要安装一次就行。

1. PyTorch 环境安装

基础环境配置参考官方文档。

2. 其他依赖安装

安装 requirements.txt 文件的环境，建议检查是否需要注释掉某些行以避免覆盖已安装的 pytorch 版本。

没有这个文件可以自己新建一个 requirements.txt，然后把下面代码复制进去：

# Ultralytics requirements
# Example: pip install -r requirements.txt
# Base ----------------------------------------
matplotlib>=3.3.0
numpy==1.24.4
opencv-python>=4.6.0
pillow>=7.1.2
pyyaml>=5.3.1
requests>=2.23.0
scipy>=1.4.1
tqdm>=4.64.0
# Logging -------------------------------------
# tensorboard>=2.13.0
# dvclive>=2.12.0
# clearml
# comet
# Plotting ------------------------------------
pandas>=1.1.4
seaborn>=0.11.0
# Export --------------------------------------
# coremltools>=7.0
# onnx>=1.12.0
# onnxsim>=0.4.1
# nvidia-pyindex
# nvidia-tensorrt
# scikit-learn==0.19.2
# tensorflow>=2.4.1
# tflite-support
# tensorflowjs>=3.9.0
# openvino-dev>=2023.0
# Extras --------------------------------------
psutil
py-cpuinfo
thop>=0.1.1
# ipython
# albumentations>=1.0.3
# pycocotools>=2.0.6
# roboflow

四、YOLOv11 训练

(1) 在根目录新建一个 python 文件，取名为：train.py。

(2) 把训练代码复制到 train.py 文件。

训练的代码如下：

# -*- coding: utf-8 -*-
import warnings
warnings.filterwarnings('ignore')
from ultralytics import YOLO

if __name__ == '__main__':
    model = YOLO(model=r'./cfg/models/11/yolo11.yaml')
    # model.load('yolo11n.pt') # 加载预训练权重，改进或者做对比实验时候不建议打开
    model.train(data=r'data.yaml', imgsz=640, epochs=50, batch=4, workers=0, device='', optimizer='SGD', close_mosaic=10, resume=False, project='runs/train', name='exp', single_cls=False, cache=False,)

注意：模型配置路径改成你自己的路径，还有数据集配置文件也修改成你自己的路径。

训练代码的参数解释：

• model 参数：该参数填入模型配置文件的路径，改进的话建议不需要填预训练模型权重
• data 参数：该参数可以填入训练数据集配置文件的路径
• imgsz 参数：该参数代表输入图像的尺寸，指定为 640x640 像素
• epochs 参数：该参数代表训练的轮数
• batch 参数：该参数代表批处理大小，电脑显存越大，就设置越大，根据自己电脑性能设置
• workers 参数：该参数代表数据加载的工作线程数，出现显存爆了的话可以设置为 0，默认是 8
• device 参数：该参数代表用哪个显卡训练，留空表示自动选择可用的 GPU 或 CPU
• optimizer 参数：该参数代表优化器类型
• close_mosaic 参数：该参数代表在多少个 epoch 后关闭 mosaic 数据增强
• resume 参数：该参数代表是否从上一次中断的训练状态继续训练。设置为 False 表示从头开始新的训练。如果设置为 True，则会加载上一次训练的模型权重和优化器状态，继续训练。
• project 参数：该参数代表项目文件夹，用于保存训练结果
• name 参数：该参数代表命名保存的结果文件夹
• single_cls 参数：该参数代表是否将所有类别视为一个类别，设置为 False 表示保留原有类别
• cache 参数：该参数代表是否缓存数据，设置为 False 表示不缓存。

一般做科研改进工作时候可以不用预训练权重，因为用预训练模型整体精度很难提高。

我这里演示加载预训练权重，训练输出如下所示：

五、YOLOv11 推理

（1）官网的预训练模型下载

进入官网的源码下载地址，往下面拉，看到模型位置，YOLOv11 针对不同的场景和应用提供了 YOLOv11n、YOLOv11s 等不同大小的模型，具体看官网提供的，需要下载哪个，鼠标左键单击下载就行。

（2）在根目录新建一个 python 文件，取名为：detect.py。

（3）把推理代码复制到 detect.py 文件。

注意：模型路径改成你自己的路径，还有预测图像也改成你自己的路径。

推理的代码如下：

# -*- coding: utf-8 -*-
from ultralytics import YOLO

if __name__ == '__main__':
    # Load a model
    model = YOLO(model=r'./yolo11n-seg.pt')
    model.predict(source=r'./assets/bus.jpg', save=True, show=True,)

推理代码的参数解释

1. model 参数：该参数可以填入模型文件路径
2. source 参数：该参数可以填入需要推理的图片或者视频路径，如果打开摄像头推理则填入 0 就行
3. save 参数：该参数填入 True，代表把推理结果保存下来，默认是不保存的，所以一般都填入 True
4. show 参数：该参数填入 True，代表把推理结果以窗口形式显示出来，默认是显示的，这个参数根据自己需求打开就行，不显示你就填 False 就行

分割模型推理结果如下：

目标检测模型推理结果如下：

六、解决训练过程中断怎么继续上次训练

在训练过程不小心中断了，那怎么继续上次的训练了，这里先不慌，官网也的代码写得非常好，它有这个断点训练功能，那么 YOLOv8 v10 v11 处理的方法都是一模一样，接下来直接看图操作就行：

• model 参数：该参数填入上次中断的模型，为 last.pt
• resume 参数：该参数设置为 True，则会加载上一次训练的模型权重和优化器状态，继续训练。

总结

YOLOv11 训练自己数据集和推理到此结束，如有问题可查阅相关文档。