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 第二处