怎樣在自定義數據集上訓練YOLOv5

怎樣在自定義數據集上訓練YOLOv5

目錄

1. 引言
2. YOLOv5簡介
3. 準備工作
   • 3.1 環境配置
   • 3.2 數據集準備
4. 數據標注
   • 4.1 標注工具
   • 4.2 標注格式
5. 數據集劃分
   • 5.1 訓練集、驗證集和測試集
   • 5.2 數據集目錄結構
6. 配置文件
   • 6.1 數據配置文件
   • 6.2 模型配置文件
7. 訓練模型
   • 7.1 訓練命令
   • 7.2 訓練參數
8. 模型評估
   • 8.1 評估指標
   • 8.2 評估命令
9. 模型推理
   • 9.1 推理命令
   • 9.2 推理結果可視化
10. 模型優化
    • 10.1 數據增強
    • 10.2 超參數調優
11. 常見問題與解決方案
12. 總結

引言

YOLOv5 是 Ultralytics 公司開發的一個高效的目標檢測模型，它在速度和精度之間取得了良好的平衡。相比于 YOLOv4，YOLOv5 在實現上更加簡潔，且易于使用。本文將詳細介紹如何在自定義數據集上訓練 YOLOv5 模型，涵蓋從數據準備到模型訓練、評估和推理的全過程。

YOLOv5簡介

YOLOv5 是基于 YOLO（You Only Look Once）系列的目標檢測模型，它通過單次前向傳播即可完成目標檢測任務。YOLOv5 的主要特點包括：

• 高效性：YOLOv5 在保持較高檢測精度的同時，具有較快的推理速度。
• 易用性：YOLOv5 提供了簡潔的代碼實現和豐富的文檔，便于用戶快速上手。
• 靈活性：YOLOv5 支持多種模型尺寸（如 YOLOv5s、YOLOv5m、YOLOv5l、YOLOv5x），用戶可以根據需求選擇合適的模型。

準備工作

3.1 環境配置

在開始訓練之前，首先需要配置好訓練環境。YOLOv5 基于 PyTorch 框架，因此需要安裝 PyTorch 及其相關依賴。

# 安裝 PyTorch
pip install torch torchvision torchaudio

# 克隆 YOLOv5 倉庫
git clone https://github.com/ultralytics/yolov5
cd yolov5

# 安裝依賴
pip install -r requirements.txt

3.2 數據集準備

在訓練 YOLOv5 之前，需要準備好自定義數據集。數據集應包含圖像和對應的標注文件。標注文件通常采用 YOLO 格式，每個圖像對應一個 .txt 文件，文件中包含目標的類別和邊界框信息。

數據標注

4.1 標注工具

常用的數據標注工具有 LabelImg、LabelMe 等。這些工具可以幫助用戶快速標注圖像中的目標，并生成對應的標注文件。

• LabelImg：支持 Pascal VOC 和 YOLO 格式的標注。
• LabelMe：支持多邊形標注，適用于不規則目標的標注。

4.2 標注格式

YOLO 格式的標注文件包含以下信息：

• 類別索引：目標的類別索引，從 0 開始。
• 邊界框坐標：歸一化的邊界框坐標，格式為 (x_center, y_center, width, height)。

例如，一個標注文件可能如下所示：

0 0.5 0.5 0.3 0.4
1 0.2 0.3 0.1 0.2

數據集劃分

5.1 訓練集、驗證集和測試集

在訓練模型之前，通常需要將數據集劃分為訓練集、驗證集和測試集。常見的劃分比例為 70% 訓練集、20% 驗證集和 10% 測試集。

5.2 數據集目錄結構

YOLOv5 要求數據集按照以下目錄結構組織：

datasets/
  custom/
    images/
      train/
      val/
      test/
    labels/
      train/
      val/
      test/

其中，images 目錄下存放圖像文件，labels 目錄下存放對應的標注文件。

配置文件

6.1 數據配置文件

YOLOv5 需要一個數據配置文件來指定數據集的路徑和類別信息。數據配置文件的格式如下：

train: datasets/custom/images/train
val: datasets/custom/images/val

nc: 2  # 類別數量
names: ['class1', 'class2']  # 類別名稱

6.2 模型配置文件

YOLOv5 提供了多種預定義的模型配置文件，用戶可以根據需求選擇合適的模型。模型配置文件的路徑為 models/yolov5s.yaml、models/yolov5m.yaml 等。

訓練模型

7.1 訓練命令

使用以下命令開始訓練模型：

python train.py --img 640 --batch 16 --epochs 100 --data datasets/custom/data.yaml --cfg models/yolov5s.yaml --weights yolov5s.pt

7.2 訓練參數

• --img：輸入圖像的尺寸。
• --batch：批量大小。
• --epochs：訓練輪數。
• --data：數據配置文件的路徑。
• --cfg：模型配置文件的路徑。
• --weights：預訓練權重的路徑。

模型評估

8.1 評估指標

常用的目標檢測評估指標包括：

• mAP（mean Average Precision）：平均精度均值，衡量模型在不同 IoU 閾值下的檢測精度。
• Precision：精確率，衡量模型預測為正樣本的樣本中實際為正樣本的比例。
• Recall：召回率，衡量模型正確檢測出的正樣本占所有正樣本的比例。

8.2 評估命令

使用以下命令評估模型：

python val.py --data datasets/custom/data.yaml --weights runs/train/exp/weights/best.pt

模型推理

9.1 推理命令

使用以下命令進行模型推理：

python detect.py --source datasets/custom/images/test --weights runs/train/exp/weights/best.pt

9.2 推理結果可視化

推理結果將保存在 runs/detect/exp 目錄下，用戶可以通過查看生成的圖像來評估模型的檢測效果。

模型優化

10.1 數據增強

數據增強是提高模型泛化能力的有效手段。YOLOv5 支持多種數據增強方法，如隨機裁剪、旋轉、翻轉等。用戶可以通過修改數據配置文件來啟用或禁用特定的數據增強方法。

10.2 超參數調優

超參數調優是模型優化的重要環節。常見的超參數包括學習率、批量大小、優化器等。用戶可以通過實驗來尋找最優的超參數組合。

常見問題與解決方案

1. 訓練過程中出現內存不足的問題：可以嘗試減小批量大小或降低圖像分辨率。
2. 模型過擬合：可以嘗試增加數據增強方法或使用正則化技術。
3. 推理速度慢：可以嘗試使用更小的模型尺寸或優化推理代碼。

總結

本文詳細介紹了如何在自定義數據集上訓練 YOLOv5 模型。通過合理的數據準備、模型訓練和優化，用戶可以構建出高效的目標檢測模型。希望本文能為讀者提供有價值的參考，幫助大家在實踐中取得更好的效果。