如何使用 Ultralytics 来训练 YOLO11 模型
前言
在 LLM 领域,我们常常能听到“大模型”这一说法:LLM 的研究人员将正常人类的自然交流对话作为语料,通过某些方式构建了一些算法,然后通过大批次、多轮次地训练来找到一组最优的参数,最终实现程序能对用户的不同输入做出不同响应的效果,仿佛有了生命一般。而在计算机视觉领域亦是同理,通过一系列图像与标签作为数据集,通过不断地训练,最终实现对图像的分类、目标检测甚至是分割功能。而在目标检测/实例分割领域,YOLO 算法自然是位于聚光灯下的主角 —— 单阶段检测算法中的王者,兼顾效率与精度。因此接下来,本文将借助 Ultralytics 库,使用自定义的数据来训练一个 YOLO11 的实例分割模型。
制作可用于 YOLO 的数据集
从 Label Studio 中导出标注任务
实际上,能为图像做标注任务的软件有很多,如 HumanSignal 的 Label Studio、LabelImg,或者是 labelme、ISAT-SAM,又或者说是 Roboflow …… 本文以 Label Studio 为例,其它的标注软件的导出也应与 Label Studio 类似。
进入需要导出的 Label Studio 项目中,在界面右上角有一个“Export”按钮,点击 Export,在弹出的菜单中选择“YOLO with Images”
很快,浏览器就会以下载的方式将标注结果导出至本地磁盘。将其解压至一个空目录中,为后面分割数据集做准备。
将数据集按照一定比例划分训练集、验证集与测试集
新建一个文件夹,以笔者为例,新建了一个名为 dataset
的文件夹,然后在 /dataset
下新建三个目录:
/train
/val
/test
这三个目录分别对应训练集、验证集与测试集。
然后,在这三个文件夹下,再分别创建两个目录:
/images
/labels
其中,/images
用于存放该集合对应的图像,/labels
用于存放对应集合的标注。
下面就是制作数据集中最重要的一步了:将原先从 Label Studio 中导出的标注数据集转化为可供 YOLO 训练验证与测试用的数据集。这一步可以用下面三个流程来概括:
- 确定训练集、验证集与测试集的比例
- 将图像与对应的
.txt
绑定,将这些绑定的个体在不超出预设数值的情况下随机地向训练集、验证集与测试集填充 - 依据 Label Studio 导出标注中的
notes.json
文件与classes.txt
文件将最终训练集中的data.yaml
文件完善
既然是随机填充,那么就不能由我们自己来手动选择哪些图片放入哪个集合当中了。因此,我们借助计算机的力量,使用编程来为每个集合随机分配图像与标签,笔者使用 Python
编写了一个简单的程序,帮助我们将图像随机拷贝至各集合对应的目录中:
import os
import shutil
import random
# 定义 Label Studio 导出数据文件夹所在路径
BASE_DIR = r'<Label Studio 导出数据文件夹所在路径,替换为自己的>'
IMAGE_DIR = os.path.join(BASE_DIR, 'images')
LABEL_DIR = os.path.join(BASE_DIR, 'labels')
# 定义需要制作成 YOLO 可用的文件夹路径
OUT_DIR = r'<需要制作成 YOLO 可用的文件夹路径,也就是 /dataset ,替换为自己的>'
TRAIN_DIR = os.path.join(OUT_DIR, 'train')
TEST_DIR = os.path.join(OUT_DIR, 'test')
VAL_DIR = os.path.join(OUT_DIR, 'val')
# 获取图像列表并打乱
images = [f for f in os.listdir(IMAGE_DIR) if f.endswith(('.jpg', '.png', '.jpeg'))]
random.shuffle(images)
# 按照 7:2:1 划分(训练集:验证集:测试集)
total = len(images)
train_size = int(total * 0.7)
val_size = int(total * 0.2)
test_size = total - train_size - val_size
# 从随机打乱的列表中切片
train_images = images[:train_size]
test_images = images[train_size:train_size + test_size]
val_images = images[train_size + test_size:]
# 复制图像至目标目录
def copy_files(file_list, split_dir):
for img in file_list:
# 移动图片
shutil.copy(os.path.join(IMAGE_DIR, img), os.path.join(split_dir, 'images', img))
# 移动对应的标签文件(若存在)
label_file = os.path.splitext(img)[0] + '.txt'
if os.path.exists(os.path.join(LABEL_DIR, label_file)):
shutil.copy(os.path.join(LABEL_DIR, label_file), os.path.join(split_dir, 'labels', label_file))
# 对三个目录分别进行图像与标签的拷贝
copy_files(train_images, TRAIN_DIR)
copy_files(test_images, TEST_DIR)
copy_files(val_images, VAL_DIR)
# 输出各集合的数量
print(f'Total images: {total}')
print(f'Train set: {train_size} images')
print(f'Test set: {test_size} images')
print(f'Validation set: {val_size} images')
最后,在 /dataset
文件夹中创建一个 YAML
文件,如 data.yaml
,并根据 Label Studio 导出数据文件夹中的 notes.json
和 classes.txt
进行编写。请遵循下方给出的格式模版进行改写:
train: # 训练集所在目录,替换为实际的,最好使用绝对路径,如/home/abc/dataset/train/
val: # 训练集所在目录,替换为实际的,最好使用绝对路径
test: # 训练集所在目录,替换为实际的,最好使用绝对路径
nc: # 标注时目标的类别数量,如 3
names: # 所有标注的类别,必须严格按照 notes.json 或 classes.txt 的顺序来填写,形如 ["Class 1", "Class 2", ... "Class n"]
# 假设 classes.txt 中的类别为:
# Apple
# Car
# Dog
# 那么此处的格式为:["Apple", "Car", "Dog"]
以笔者为例,将根据 notes.json
对 data.yaml
做如下编辑:
由 Label Studio 导出标注数据目录中的 notes.json
:
至此,我们完成了用于 YOLO 训练使用的数据集,接下来就可以使用这个数据集进行模型的训练了。
训练 YOLO11 模型
在另一处目录创建一个新的 Python
文件,使用任意一个 IDE(此处使用 VSCode) 打开此文件夹并在该 Python
文件中输入以下内容:
# 从 Ultralytics 模块中导入 YOLO 类
from ultralytics import YOLO
# 创建 YOLO 对象
model = YOLO('<需要使用的预训练模型,比如 yolo11m.pt>')
# 使用 YOLO 对象的 train 方法进行训练
model.train(data= # 自定义的 YOLO 训练用数据集的 YAML 配置文件路径,结尾记得加英语逗号“,”
epochs= # 要训练的轮次,比如 500
patience= # 耐心度,比如 300,若在 300 轮后没有任何进步,则会中断训练以节约开销
batch= # 训练批次,若硬件性能越好,则数量可以调的越高。一般 4-16 为普遍值,一般为 2 的倍数
imgsz= # 输入的特征图大小,在训练时只能是一个 32 倍数的整数(方阵),如 720、960 或 1280
device= # 用于训练的设备,0 代表第 1 张 GPU,'cpu' 则代表使用 CPU 进行计算
)
重要提醒:从现在开始,请保持终端中的工作路径位于该 Python
文件所在的目录,不要变更,否则会出现非常多难以预料的错误。比如说目前 Python
文件位于 ~/VSCode Projects/train_yolo11
,请保证终端中的工作路径位于 ~/VSCode Projects/train_yolo11
,不得变换到如 ~
的其它位置。若工作目录位于其它位置,请使用 cd
命令切换回当前目录下。
此外,若您使用了虚拟环境,请将 Python
解释器设置为安装了 Ultralytics 的解释器。以笔者为例,使用 Miniconda 创建了一个名为 yolo11_devel
的环境,则需要在 VSCode 中将解释器设置为 3.9.21 ('yolo11_devel')
。
需要了解更多关于预训练模型及模型的训练参数,请参阅 Ultralytics 训练文档 。
请根据实际情况对上述内容进行修改,以笔者为例,使用 yolo11m-seg.pt
预训练模型,训练 500 轮,在 400 轮后允许自断,每一批次训练量设为 4,输入特征图由 1920 × 1080 缩放为 960 × 960,使用我的第一块 GPU 进行训练,因此笔者的 Python
训练代码如下图所示:
运行该 Python
文件,我们就可以在终端中看到,程序会自动从 Ultralytics 下载必要的一些权重文件与资源文件,在加载和预处理完毕后,模型就会开始训练:
笔者的 GPU 为 NVIDIA GeForce RTX 2070 SUPER 8G,在 batch 设置为 4 的情况下,平均训练 1 轮大约花费 30-35 秒。按照 35 秒一轮的速度,500 轮的训练大约需要 4-5 个小时。
历时近 5 个小时的训练完成后,我们会在当前 Python
文件的目录的 /runs/***/train/
下获得训练的结果,/runs/***/train/
下的各类图像给出了训练/验证中获得的损失函数、混淆矩阵,$mAP$ 等结果,而 /runs/***/train/weights/
中则包含了两个最终获得的权重模型 best.pt
和 last.pt
,分别代表最佳的模型与最后一轮获得的模型。
至此,我们已成功训练并获得了自定义数据集的 YOLO11 实例分割模型。实际上,目标检测模型的训练流程与本文所示范的实例分割训练流程基本相同,感兴趣的读者可以自行尝试 Ultralytics 8.3 的目标检测模型训练,且训练开销相较于实例分割会低不少。
测试 YOLO11 模型
复制 /runs/***/train/weights/
中权重文件的绝对路径,如在 VSCode 中复制 best.pt
的绝对路径,然后在终端中,使用以下命令来进行模型的推理测试:
# [model-path] 为刚才复制的权重文件绝对路径
# [media-to-predict-path] 为需要推理的媒体文件,可以是图像也可以是视频
yolo segment predict model=[model-path] source=[media-to-predict-path]
以笔者为例,推理命令如下图所示:
按下 Enter 运行,即可看到终端提示在当前目录下的 /runs/segment/predict/
下产生了推理结果,我们打开该目录,即可查看到推理后的图像:
非常完美,图中的几个辣椒目标都被正确地识别并分割了出来,类别也没有识别成其它颜色(如把红色识别成绿色)。大功告成,接下来就可以使用该模型来进行其它工程项目的开发与应用了。
总结
本文借助 Ultralytics 演示了如何使用自定义数据集来训练一个 YOLO11 的实例分割模型,并在此过程中演示了手动划分数据集、训练参数的调整与模型推理测试。实际上,笔者往期都是使用如 Roboflow 这样的在线网站直接导出可供 YOLO 训练用的数据集,然后在本地进行目标检测模型的训练,而这一部分内容迟迟(Ultralytics - THU-MIG YOLOv10 的目标检测模型的训练与推理测试)未写出的原因正是笔者注意到了 Ultralytics 从 8.1 开始对实例分割/定向目标检测的支持度在不断完善与提升,但是 Roboflow 等在线网站却对数据集的标注与数量开始有了严格的限制…… 最后,如果各位对 Ultralytics YOLO11 的模型训练流程有其他见解或想说的,欢迎在下方评论区留言。