基于 YOLO V10 Fine-Tuning 训练自定义的目标检测模型

一、YOLO V10

在本专栏的前面几篇文章中，我们使用

ultralytics

公司开源发布的

YOLO-V8

模型，分别

Fine-Tuning

实验了目标检测、关键点检测、分类任务，实验后发现效果都非常的不错，但它已经不是最强的了。最新的

YOLO-V10

已经完全超越之前的所有版本，

YOLO-V10

由清华大学提供，采用无

NMS

训练和效率-精度驱动架构，提供目前最先进的性能和延迟。

在这里插入图片描述

从上图中的对比效果可以明显看出，

YOLO-V10

不仅在速度上得到了极大的提升，精度同样也得到了明显的提升。主要得益于其无

NMS

训练的重大变化。

在模型上

V10

和之前的版本类似，包括不同大小的模型，从小到大包括：

YOLOv10-N：用于资源极其有限环境的纳米版本。
YOLOv10-S：兼顾速度和精度的小型版本。
YOLOv10-M：通用中型版本。
YOLOv10-B：平衡型，宽度增加，精度更高。
YOLOv10-L：大型版本，精度更高，但计算资源增加。
YOLOv10-X：超大型版本可实现最高精度和性能。

模型的比较如下：

在这里插入图片描述

更多的介绍可以参考官方的文档：

https://docs.ultralytics.com/de/models/yolov10/#model-variants

本文借助

ultralytics

框架对

YOLO V10

迁移训练自定义的目标检测模型，本次实验训练一个人脸检测模型，包括数据标注、数据拆分、训练、测试等过程。

实验采用

ultralytics

框架，可以帮助开发人员高效完成数据训练和验证任务，由于

ultralytics

默认采用的为

PyTorch

框架，因此实验前请安装好

cuda

和

torch

环境，如果没有

GPU

环境，由于

YOLO V10

已经足够轻量级，使用

CPU

也是可以训练。

安装

ultralytics

库：

pip install ultralytics -i https://pypi.tuna.tsinghua.edu.cn/simple

如果已经安装，需要更新到最新版本：

pip install--upgrade ultralytics -i https://pypi.tuna.tsinghua.edu.cn/simple

ultralytics

使用文档：

https://docs.ultralytics.com/zh/quickstart/#use-ultralytics-with-python

测试

YOLO V10

的效果：

测试图片：
在这里插入图片描述

这里使用

yolov10n

模型，如果模型不存在会自动下载

from ultralytics import YOLO
# Load a model
model = YOLO('yolov10n.pt')
results = model.predict('./img/1.png')
results[0].show()

在这里插入图片描述

二、准备训练数据及标注

图像数据可以从网上找一些或者自己拍摄，我这里准备了一些人的图片：

在这里插入图片描述

这里可以准备两个目录，

data/images

和

data/labels

，其中

labels

存放标注后的文件，将收集到的图像放在

images

目录下：

在这里插入图片描述

下面使用

labelimg

工具进行标注，如果没有安装，使用下面命令安装：

pip install labelimg -i https://pypi.tuna.tsinghua.edu.cn/simple

然后在控制台输入：

labelimg

打开可视化工具：

在这里插入图片描述

注意：数据集格式默认是

VOC

格式的，要选择为

YOLO

，我这里的人脸标签为

face

，这个后面需要使用到。

标注完成后，可以在

/data/labels

下看到标注后的文件：

在这里插入图片描述

三、数据拆分

这里拆分为

90%

的训练集，

10%

的验证集，这部分和之前训练

YOLO V8

时一致，拆分脚本如下，

import os
import shutil
from tqdm import tqdm
# 图片地址
image_dir ="data/images/"# 标准文件地址
label_dir ="data/labels/"# 训练集的比例
training_ratio =0.9# 拆分后数据的位置
train_dir ="train_data"defsplit_data():list= os.listdir(image_dir)all=len(list)
    train_count =int(all* training_ratio)
    train_images =list[0:train_count]
    val_images =list[train_count:]# 训练集目录
    os.makedirs(os.path.join(train_dir,"images/train"), exist_ok=True)
    os.makedirs(os.path.join(train_dir,"labels/train"), exist_ok=True)# 验证集目录
    os.makedirs(os.path.join(train_dir,"images/val"), exist_ok=True)
    os.makedirs(os.path.join(train_dir,"labels/val"), exist_ok=True)# 训练集withopen(os.path.join(train_dir,"train.txt"),"w")asfile:file.write("\n".join([train_dir +"images/train/"+ image_file for image_file in train_images]))print("save train.txt success!")# 拷贝数据for item in tqdm(train_images):
        label_file = item.replace(".jpg",".txt")
        shutil.copy(os.path.join(image_dir, item), os.path.join(train_dir,"images/train/"))
        shutil.copy(os.path.join(label_dir, label_file), os.path.join(train_dir,"labels/train/"))# 验证集withopen(os.path.join(train_dir,"val.txt"),"w")asfile:file.write("\n".join([train_dir +"images/val/"+ image_file for image_file in val_images]))print("save val.txt success!")# 拷贝数据for item in tqdm(val_images):
        label_file = item.replace(".jpg",".txt")
        shutil.copy(os.path.join(image_dir, item), os.path.join(train_dir,"images/val/"))
        shutil.copy(os.path.join(label_dir, label_file), os.path.join(train_dir,"labels/val/"))if __name__ =='__main__':
    split_data()

在这里插入图片描述
可以在

train_data

中看到拆分后的数据集格式：

在这里插入图片描述

四、训练

使用

ultralytics

框架训练非常简单，仅需三行代码即可完成训练，不过在训练前需要编写

YAML

配置信息，主要标记数据集的位置。

创建

face.yaml

文件，写入下面内容：

path: D:/pyProject/yolov10/train_data # 数据集的根目录, 建议使用绝对路径train: images/train # 训练集图像目录val: images/val # 验证集图像目录test:# test images (optional)# 分类names:0: face

注意分类中的

face

就是上面标注时的标签名。

开始训练：

from ultralytics import YOLO
# 加载模型
model = YOLO('yolov10n.pt')# 训练
model.train(
    data='face.yaml',# 训练配置文件
    epochs=100,# 训练的周期
    imgsz=640,# 图像的大小
    device=[0],# 设备，如果是 cpu 则是 device='cpu'
    workers=0,
    lr0=0.0001,# 学习率
    batch=8,# 批次大小
    amp=False# 是否启用混合精度训练)

运行后可以看到打印的网络结构：

在这里插入图片描述

训练中：

在这里插入图片描述

训练结束后可以在

runs

目录下面看到训练的结果：

在这里插入图片描述

其中

weights

下面的就是训练后保存的模型，这里可以先看下训练时

loss

的变化图：

在这里插入图片描述

五、模型测试

在

runs\detect\train\weights

下可以看到

best.pt

和

last.pt

两个模型，表示最佳和最终模型，下面使用

best.pt

模型进行测试

from ultralytics import YOLO
from matplotlib import pyplot as plt
import os
plt.rcParams['font.sans-serif']=['SimHei']# 测试图片地址
base_path ="test"# 加载模型
model = YOLO('runs/detect/train/weights/best.pt')for img_name in os.listdir(base_path):
    img_path = os.path.join(base_path, img_name)
    image = plt.imread(img_path)# 预测
    results = model.predict(image, device='cpu')
    boxes = results[0].boxes.xyxy
    confs = results[0].boxes.conf
    ax = plt.gca()for index, boxe inenumerate(boxes):
        x1, y1, x2, y2 = boxe[0], boxe[1], boxe[2], boxe[3]
        score = confs[index].item()
        ax.add_patch(plt.Rectangle((x1, y1),(x2 - x1),(y2 - y1), linewidth=2, fill=False, color='red'))
        plt.text(x=x1, y=y1-10, s="{:.2f}".format(score), fontsize=15, color='white',
                 bbox=dict(facecolor='black', alpha=0.5))
    plt.imshow(image)
    plt.show()