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

一、YOLO V8YOLO V8是由2023年ultralytics公司开源的发布是结合了前几代YOLO的融合改进版。YOLO V8支持全方位的视觉AI任务包括检测、分割、姿态估计、跟踪和分类。并且在速度和准确性方面具有无与伦比的性能。能够应用在各种对速度和精度要求较高的应用和领域。网络结构如下图所示YOLO V8相对于YOLO V5还是有很大的不同例如YOLO V8相对于YOLO V5依然使用的是CSP的思想不过将V5中的C3模块换成了C2F模块以减轻模型的大小也依旧使用V5架构中的SPPF模块。但是在PAN-FPN层面V8将V5中的上采样阶段中的卷积结构去除了。同时借鉴了YOLOX的Decoupled-Head结构分类和回归两个任务的HEAD不再共享参数等。在模型上 V8 和 V5 类似包括不同大小的模型从小到大包括yolov8n、yolov8s、yolov8m、yolov8l、yolov8x等模型的比较如下更多的介绍可以参考官方的文档https://docs.ultralytics.com/de/models/yolov8/本文借助ultralytics中YOLO V8迁移训练自定义的目标检测模型在本次的实验中主要训练一个人脸检测模型包括数据标注、数据拆分、训练、测试等过程。本次采用ultralytics公司发布的ultralytics框架可以帮助开发人员高效完成数据训练和验证任务由于ultralytics默认采用的为PyTorch框架因此实验前请安装好cuda和torch环境如果没有GPU环境由于YOLO V8已经足够轻量级使用CPU也是可以训练。安装ultralytics库pipinstallultralytics-ihttps://pypi.tuna.tsinghua.edu.cn/simpleultralytics使用文档https://docs.ultralytics.com/zh/quickstart/#use-ultralytics-with-python测试YOLO V8的效果测试图片这里使用yolov8n模型如果模型不存在会自动下载fromultralyticsimportYOLO# Load a modelmodelYOLO(yolov8n.pt)# pretrained YOLOv8n modelresultsmodel.predict(./img/1.png)# Show resultsresults[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%的验证集拆分脚本如下importosimportshutilfromtqdmimporttqdm# 图片地址image_dirdata/images/# 标准文件地址label_dirdata/labels/# 训练集的比例training_ratio0.9# 拆分后数据的位置train_dirtrain_datadefsplit_data():listos.listdir(image_dir)alllen(list)train_countint(all*training_ratio)train_imageslist[0:train_count]val_imageslist[train_count:]# 训练集目录os.makedirs(os.path.join(train_dir,images/train),exist_okTrue)os.makedirs(os.path.join(train_dir,labels/train),exist_okTrue)# 验证集目录os.makedirs(os.path.join(train_dir,images/val),exist_okTrue)os.makedirs(os.path.join(train_dir,labels/val),exist_okTrue)# 训练集withopen(os.path.join(train_dir,train.txt),w)asfile:file.write(\n.join([train_dirimages/train/image_fileforimage_fileintrain_images]))print(save train.txt success!)# 拷贝数据foritemintqdm(train_images):label_fileitem.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_dirimages/val/image_fileforimage_fileinval_images]))print(save val.txt success!)# 拷贝数据foritemintqdm(val_images):label_fileitem.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/yolov8/train_data# 数据集的根目录, 建议使用绝对路径train:images/train# 训练集图像目录val:images/val# 验证集图像目录test:# test images (optional)# 分类names:0:face注意分类中的face就是上面标注时的标签名。开始训练fromultralyticsimportYOLO# 加载模型modelYOLO(yolov8n.pt)# 训练model.train(dataface.yaml,# 训练配置文件epochs50,# 训练的周期imgsz640,# 图像的大小device[0],# 设备如果是 cpu 则是 devicecpuworkers0,lr00.001,# 学习率batch8,# 批次大小ampFalse# 是否启用混合精度训练)运行后可以看到打印的网络结构训练中训练结束后可以在runs目录下面看到训练的结果其中weights下面的就是训练后保存的模型这里可以先看下训练时loss的变化图五、模型测试使用best.pt模型fromultralyticsimportYOLO# 加载模型modelYOLO(runs/detect/train/weights/best.pt)# 预测resultsmodel.predict(data/images/8.jpg)# Show resultsresults[0].show()