YOLOX 訓練自定義數據 - VOC2007數據集

官方文檔 - Train Custom Data

環境說明

• Windows11
• Git，可選
• Python 3.10.9 [下載地址]
• GTX 1050Ti
• CUDA Toolkits 12

先決條件

• 安裝 Nvidia 顯卡驅動 [下載地址]
• 安裝 Nvidia CUDA Toolkits [下載地址]

主要步驟

• 一、下載 YOLOX 源代碼
• 二、按需修改 YOLOX 源代碼
• 三、製作 VOC2007 數據集
• 四、訓練
• 五、預測

一、下載 YOLOX 源代碼

1.1 克隆源代碼

克隆源代碼到指定目錄，並進入YOLOX目錄：

git clone https://github.com/Megvii-BaseDetection/YOLOX.git
cd ./YOLOX

如未安裝 Git 工具，也可以在 GitHub 上直接下載壓縮包。

1.2 創建虛擬環境

在根目錄創建 venv 虛擬環境，並激活：

python -m venv .venv
.venv/Scripts/activate

注意：後續所有命令的運行，凡是需要運行 Python 的，都需要在已激活虛擬環境的終端中運行。

1.3 安裝依賴

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

# 以開發模式安裝 yolox 自身
python -m pip install -v -e ./

使用 pip 的鏡像地址可以加快下載速度，如清華大學鏡像。

在 Windows 的用戶目錄下，創建 pip/pip.ini 文件並寫入以下內容即可：

[global]
index-url=https://pypi.tuna.tsinghua.edu.cn/simple

1.4 安裝 PyTorch 的 GPU 版本

# 先卸載 CPU 版本
python -m pip uninstall torch torchvision torchaudio

# 再安裝 GPU 版本
python -m pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

參考 PyTorch 官網

1.5 驗證

官方文檔 - Get Started

根據官方文檔說明，需要先下載 YOLOX 預訓練好的權重文件，這裏以 yolox_s.pth 爲例：

[前往下載頁面]

可以將權重文件下載到任意目錄，但爲了方便，我們選擇下載到 YOLOX 的 weights/ 目錄，因此需要創建 weights/ 目錄。

然後在根目錄下運行：

# PowerShell

python tools/demo.py image `
  -n yolox-s `
  -c /path/to/your/yolox_s.pth `
  --path assets/dog.jpg `
  --conf 0.25 `
  --nms 0.45 `
  --tsize 640 `
  --save_result `
  --device [cpu/gpu]

這裏選擇了圖片 assets/dog.jpg：

如果一切正常，那麼將會在產生一個新的目錄 YOLOX_outputs/，其中的子目錄 vis_res/ 中的圖片是對 assets/dog.jpg 的預測結果圖片：

二、按需修改 YOLOX 源代碼

2.1 yolox/data/datasets/voc.py

• 第一處 大約在138行，將 / 修改爲 \\。 否則會報錯超出索引邊界。

# 修改前
path_filename = [
    (self._imgpath % self.ids[i]).split(self.root + "/")[1]
    for i in range(self.num_imgs)
]

# 修改後
path_filename = [
    (self._imgpath % self.ids[i]).split(self.root + "\\")[1]
    for i in range(self.num_imgs)
]

• 第二處 大約在264行，將 if dets == [] 改爲 if dets.size == 0。 否則報錯。

# 修改前
if dets == []:
    continue

# 修改後
if dets.size == 0:
    continue

• 第三處 大約在281行，將 "{:s}.xml" 改爲 "{}.xml"。 否則路徑拼接錯誤。

# 修改前
annopath = os.path.join(rootpath, "Annotations", "{:s}.xml")

# 修改後
annopath = os.path.join(rootpath, "Annotations", "{}.xml")

2.2 yolox/data/datasets/voc_classes.py

將 VOC_CLASSES 更改爲自定義數據集中的實際類別。

2.3 exps/example/yolox_voc/yolox_voc_s.py

• 第一處 將 self.num_classes = 20 修改爲實際的類別數量。

• 第二處 由於我們的數據集只有 VOC2007，因此修改 get_dataset 函數。

# 修改前
image_sets=[('2007', 'trainval'), ('2012', 'trainval')],

# 修改後
image_sets=[('2007', 'trainval')],

2.4 tools/demo.py

由於我們的數據集是 VOC2007，因此修改

# 修改前
from yolox.data.datasets import COCO_CLASSES

# 修改後
from yolox.data.datasets.voc_classes import VOC_CLASSES

同時將本文件中使用了 COCO_CLASSES 的其他代碼改爲 VOC_CLASSES。

三、製作 VOC2007 數據集

VOC2007 數據集的官方示例數據如下：

• 訓練和驗證數據集 - VOCtrainval_06-Nov-2007.tar
• 測試數據集 - VOCtest_06-Nov-2007.tar

其內部數據目錄結構如下：

2.1 數據標註

標註工具有：labelImg，labelMe 等，本文使用的是 labelImg。 具體用法，參考其官方文檔。

2.2 創建數據集目錄

首先在 YOLOX 源代碼 datasets/ 目錄中創建上圖中的文件結構。 接下來：

• 將圖片放入 JPEGImages 目錄
• 將標註放入 Annotations 目錄

2.3 劃分數據集

圖片數據集需要被劃分爲”訓練和驗證集“、”測試集“。

分別對應 ImageSets 目錄：

• 訓練集：train.txt
• 驗證集：val.txt
• 訓練和驗證集：trainval.txt
• 測試集：test.txt

VOC2007 的官方示例數據集的劃分：

The data has been split into 50% for training/validation and 50% for testing. The distributions of images and objects by class are approximately equal across the training/validation and test sets. In total there are 9,963 images, containing 24,640 annotated objects.

數據已分爲 50% 用於訓練/驗證，50% 用於測試。 在訓練/驗證和測試集中，按類別劃分的圖像和對象的分佈大致相等。 總共有 9,963 張圖像，包含 24,640 個帶註釋的對象。

但也可以自由設置劃分比例，因爲需要自行編寫代碼來劃分。

劃分數據集可以自行編寫代碼來實現，也可以藉助現有庫來實現，如 sklearn 的 sklearn.model_selection.train_test_split 函數。

以下是用於劃分的 Python 代碼示例：

爲了便於運行，將示例代碼移動到 datasets/voc_split.py。

import os
import random


def dataset_path(relative_path):
    return os.path.abspath(
        os.path.join(os.path.dirname(__file__), os.path.normpath(relative_path))
    )


def name(filename):
    basename = os.path.basename(filename)
    return os.path.splitext(basename)[0]


def split(voc_path, train_ratio=0.5, val_ratio=0.2):
    jpeg_images_folder = os.path.join(voc_path, "JPEGImages")

    # 博客檢測到下面這行代碼有敏感詞，請自行修正下面這行代碼。
    jpeg_images_files = os.listdr(jpeg_images_folder) 

    # 獲取文件名
    names = [os.path.splitext(file)[0] for file in jpeg_images_files]
    num = len(names)
    
    # 隨機打亂文件列表
    random.shuffle(names)
    
    # 計算劃分的索引位置
    train_split = int(train_ratio * num)
    val_split = int((train_ratio + val_ratio) * num)
    
    # 獲取劃分後的文件列表
    train = sorted(names[:train_split])
    val = sorted(names[train_split:val_split])
    test = sorted(names[val_split:])
    return { "train": train, "val": val, "test": test,  "trainval": train + val }


def write(dict, voc_path):
    image_sets_folder = os.path.join(voc_path, "ImageSets", "Main")
    os.makedirs(image_sets_folder, exist_ok=True)
    for filname, names in dict.items():
        with open(os.path.join(image_sets_folder, filname + ".txt"), "w") as f:
            for name in names:
                f.write(name + "\n")


def main():
    voc_path = dataset_path("VOCdevkit/VOC2007")
    dict = split(voc_path, 0.5, 0.2)
    write(dict, voc_path)
    print("Done")


if __name__ == "__main__":
    main()

在終端運行：

python datasets/voc_split.py

即可將同目錄下的 VOC 數據集按照給定比例劃分。

四、訓練

在終端運行下面的命令，開始訓練：

# PowerShell

python tools/train.py `
    --exp_file exps/example/yolox_voc/yolox_voc_s.py `
    --devices 1 `
    --batch-size 8 `
    --fp16 `
    --occupy `
    --ckpt weights/yolox_s.pth

最終會在 YOLOX_outputs/yolox_voc_s/ 目錄下輸出權重文件， 一般會選擇 best_ckpt.pth 作爲訓練結果。

tools/train.py 各參數的具體含義，在 tool/train.py 中有聲明。

五、預測

在終端運行下面的命令，開始預測：

# PowerShell

python tools/demo.py image `
    --exp_file exps/example/yolox_voc/yolox_voc_s.py `
    --ckpt YOLOX_outputs/yolox_voc_s/best_ckpt.pth `
    --path datasets/VOCdevkit/VOC2007/JPEGImages `
    --conf 0.3 `
    --nms 0.3 `
    --save_result `
    --device cpu

參數 --path 表示我們要預測整個 datasets/VOCdevkit/VOC2007/JPEGImages 目錄中的圖片。

如果一切正常，那麼將會在 YOLOX_outputs/ 的子目錄 vis_res/ 中輸出預測結果圖片。

遇到的問題

問題1

IndexError: list index out of range

參考： 2.1 yolox/data/datasets/voc.py 第一處

問題2

ValueError: operands could not be broadcast together with shapes (6,5) (0,)

參考： 2.1 yolox/data/datasets/voc.py 第二處