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

一、YOLO V8

YOLO 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 库：

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

ultralytics 使用文档：

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

测试 YOLO V8 的效果：

测试图片：

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

from ultralytics import YOLO
# Load a model
model = YOLO('yolov8n.pt')  # pretrained YOLOv8n model

results = model.predict('./img/1.png')
# Show results
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% 的验证集，拆分脚本如下，

import os
import shutil
from tqdm import tqdm

# 图片地址
image_dir = "data/images/"
# 标准文件地址
label_dir = "data/labels/"
# 训练集的比例
training_ratio = 0.9
# 拆分后数据的位置
train_dir = "train_data"

def split_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)

    # 训练集
    with open(os.path.join(train_dir, "train.txt"), "w") as file:
        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/"))

    # 验证集
    with open(os.path.join(train_dir, "val.txt"), "w") as file:
        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/yolov8/train_data # 数据集的根目录, 建议使用绝对路径
train: images/train # 训练集图像目录
val: images/val # 验证集图像目录
test: # test images (optional)

# 分类
names:
  0: face

注意分类中的 face 就是上面标注时的标签名。

开始训练：

from ultralytics import YOLO

# 加载模型
model = YOLO('yolov8n.pt')

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

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

训练中：

训练结束后可以在 runs 目录下面看到训练的结果：

其中 weights 下面的就是训练后保存的模型，这里可以先看下训练时 loss 的变化图：

五、模型测试

使用 best.pt 模型

from ultralytics import YOLO
# 加载模型
model = YOLO('runs/detect/train/weights/best.pt')

# 预测
results = model.predict('data/images/8.jpg')

# Show results
results[0].show()