尝鲜使用 YOLO V11 Fine-Tuning 训练自定义的目标检测模型

news2024/9/30 18:53:17

一、YOLO V11

2024年9月30日，Ultralytics官方团队宣布YOLOv11正式发布，标志着YOLO系列实时目标检测器的又一次重大升级。这一新版本不仅在准确性和检测速度上再创新高，还通过架构和训练方法的革新，极大地提升了目标检测的综合性能。

在这里插入图片描述

YOLOv11 在 YOLOv8 基础上进行了重要改进，同样支持全方位的视觉 AI 任务，包括检测、分割、姿态估计、跟踪和分类。它继续采用了卷积神经网络设计，充分利用了最新的计算能力，以提高检测效率和降低延迟。在实际应用中，无论是物体识别、跟踪还是图像分割，都展现出了卓越的性能。

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

YOLO11n：用于资源极其有限环境的纳米版本。
YOLO10s：兼顾速度和精度的小型版本。
YOLO10m：通用中型版本。
YOLO10l：大型版本，精度更高，但计算资源增加。
YOLO10x：超大型版本可实现最高精度和性能。

模型的比较如下：

在这里插入图片描述

官方的文档：

https://docs.ultralytics.com/de/models/yolo11/#supported-tasks-and-modes

使用示例：

安装 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

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

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

from ultralytics import YOLO
# Load a model
model = YOLO('yolo11n.pt')

results = model.predict('./img/1.png')
results[0].show()

在这里插入图片描述

二、微调训练

数据集使用本专栏前面实验 YOLO-V10 时标注的人脸数据集，这里你可以收集一些自定义的图片，然后根据下面文章中介绍的方式进行标注：

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

微调训练，其中 face.yaml 文件内容和上面文章 YOLO-V10 时的一致：

from ultralytics import YOLO

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

# 训练
model.train(
    data='face.yaml', # 训练配置文件
    epochs=50, # 训练的周期
    imgsz=640, # 图像的大小
    device=[0], # 设备，如果是 cpu 则是 device='cpu'
    workers=0,
    # lr0=0.01, # 学习率
    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 in enumerate(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()