从零开始训练YOLO模型

从零开始训练YOLO模型

引言：

YOLO（You Only Look Once）是一种流行的目标检测算法，以其快速和准确的特点而闻名。本文将引导你从零开始训练一个YOLO模型，包括环境准备、数据集准备、模型训练和测试等步骤。

1. 准备环境

1.1 安装CUDA

  1. 访问NVIDIA官网下载对应显卡驱动版本的CUDA Toolkit。
  2. 配置环境变量：
    • CUDA_HOME: CUDA的安装路径（例如：C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.0）
    • Path: 将%CUDA_HOME%\bin和%CUDA_HOME%\libnvvp添加到Path环境变量中。
  3. 验证CUDA是否安装成功：在命令行中输入nvcc -V，如果显示CUDA版本信息，则表示安装成功。

1.2 安装PyTorch

  1. 访问PyTorch官网，选择合适的PyTorch版本和安装方式。
    注意：
    • PyTorch版本需要与CUDA版本匹配
    • 可以使用nvidia-smi命令查看CUDA版本
    • PyTorch版本应低于CUDA版本
    • 推荐使用pip命令安装PyTorch

1.3 安装YOLOv11

  1. 使用pip命令安装YOLOv11： 
    1
    pip install ultralytics
    注意： 建议使用国内镜像源加速下载。

2. 准备工作

2.1 硬件环境

一台配备NVIDIA显卡的电脑（也可以使用AMD显卡），建议显存容量不小于4GB。

2.2 数据集准备

  1. 数据集格式： YOLO模型常用的数据集格式包括VOC格式和COCO格式。
    • VOC格式： 图片和XML标签文件分别存储。
    • COCO格式： 所有标注信息存储在一个JSON文件中。
  2. 标注工具： 常用的目标检测标注工具包括LabelImg、精灵标注等。
    • LabelImg： 简单易用，支持VOC格式。
    • 精灵标注： 功能强大，支持多种格式。
  3. 数据集划分： 将数据集划分为训练集、验证集和测试集。
    • 训练集： 用于训练模型。
    • 验证集： 用于调整模型参数。
    • 测试集： 用于评估模型性能。

2.3 模型选择

选择一个合适的YOLO模型作为起始模型。
本文选择YOLOv11作为示例。

3. 训练模型

3.1 准备数据集

新建训练项目的完整文件夹结构为：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
project_folder/
├── data/
│   ├── train/
│   │   ├── image1.jpg
│   │   ├── image2.jpg
│   │   └── ...
│   ├── val/
│   │   ├── image1.jpg
│   │   ├── image2.jpg
│   │   └── ...
│   ├── test/
│   │   ├── image1.jpg
│   │   ├── image2.jpg
│   │   └── ...
│   └── data.yaml
├── models/
│   └── yolov11.pt
└── train.py

其中train文件夹包含训练集图片，val文件夹包含验证集图片，test文件夹包含测试集图片，data.yaml是数据集配置文件，models文件夹包含模型文件，train.py是训练脚本。

data.yaml文件用于配置数据集路径、类别数和类别名称。

1
2
3
4
5
train: path/to/train/images
val: path/to/val/images
test: path/to/test/images
nc: #number_of_classes
names: ['class1', 'class2', 'class3', ...]

其中path/to/train/images是训练集图片的路径，path/to/val/images是验证集图片的路径，path/to/test/images是测试集图片的路径，nc是类别数，names是类别名称。

3.2 准备模型

下载YOLOv11的模型，并放在一个文件夹(models)中。

3.3 训练模型

  1. 修改train.py脚本：

    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    from ultralytics import YOLO

    # 加载预训练模型
    model = YOLO('models/yolov11.pt')

    # 训练模型
    model.train(
        data="data/data.yaml",
        epochs=100,
        batch=16,
        imgsz=640
    )
    • data: 数据集配置文件路径。
    • epochs: 训练轮数。
    • batch: 批处理大小。
    • imgsz: 输入图片大小。

  2. 运行train.py脚本：

    1
    python train.py

  3. 监控训练过程：

    • 使用TensorBoard等工具监控训练过程。
    • 观察loss曲线、mAP等指标。

训练完成后，模型会保存在runs/train/exp/weights/文件夹中。

4. 测试模型

4.1 准备测试数据集

将测试集图片放在一个文件夹中。

4.2 测试模型

  1. 修改test.py脚本：

    1
    2
    3
    4
    5
    6
    7
    from ultralytics import YOLO

    # 加载训练好的模型
    model = YOLO('runs/train/exp/weights/best.pt')

    # 测试模型
    results = model(test="test/images")

  2. 运行test.py脚本：

    1
    python test.py

  3. 评估模型性能：

    • 常用的目标检测测试指标包括Precision、Recall、mAP等。
    • 可以使用ultralytics提供的工具计算这些指标。