基于 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库：

pipinstallultralytics-ihttps://pypi.tuna.tsinghua.edu.cn/simple

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

pipinstall--upgradeultralytics-ihttps://pypi.tuna.tsinghua.edu.cn/simple

ultralytics使用文档：

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

测试YOLO V10的效果：

测试图片：

这里使用yolov10n模型，如果模型不存在会自动下载

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

二、准备训练数据及标注

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

这里可以准备两个目录，data/images和data/labels，其中labels存放标注后的文件，将收集到的图像放在images目录下：

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

pipinstalllabelimg-ihttps://pypi.tuna.tsinghua.edu.cn/simple

然后在控制台输入：labelimg打开可视化工具：

注意：数据集格式默认是VOC格式的，要选择为YOLO，我这里的人脸标签为face，这个后面需要使用到。

标注完成后，可以在/data/labels下看到标注后的文件：

三、数据拆分

这里拆分为90%的训练集，10%的验证集，这部分和之前训练YOLO V8时一致，拆分脚本如下，

importosimportshutilfromtqdmimporttqdm# 图片地址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_fileforimage_fileintrain_images]))print("save train.txt success!")# 拷贝数据foritemintqdm(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_fileforimage_fileinval_images]))print("save val.txt success!")# 拷贝数据foritemintqdm(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就是上面标注时的标签名。

开始训练：

fromultralyticsimportYOLO# 加载模型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模型进行测试

fromultralyticsimportYOLOfrommatplotlibimportpyplotaspltimportos plt.rcParams['font.sans-serif']=['SimHei']# 测试图片地址base_path="test"# 加载模型model=YOLO('runs/detect/train/weights/best.pt')forimg_nameinos.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()forindex,boxeinenumerate(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()