[Python知识库]YOLOX自定义数据集训练（抢先踩坑）

序言

昨天被YOLOX刷屏了，各大公众号强推：性能超yolov5！！吊打一切yolo！！看麻了我，标题还能再夸张点嘛？出于对前沿技术的渴望，还是要去学习学习，论文中改进了很多地方，这里就不再介绍，刷屏的文章中都已经介绍了，直接来看下作者给的性能对比图：

左边是大模型的性能对比，右边是各自轻量型模型的对比。可以看到越靠近左上角性能越优。此外作者还一次性开源了tensorrt、ncnn、openvino、onnx部署代码，可谓业界良心，从训练到部署一条龙服务。

习惯了yolo格式训练方式的我，在踩了无数的坑后，决定写一篇训练教程文章，记录自己踩坑的过程，也给后来人做参考。

话说在前头，坑是真的多，给我整麻了，要训练的同学做好debug的心理准备。

YOLOX官方仓库

一、配置环境

训练之前，按照官方的环境配一下，先介绍我的硬件情况：

• ubuntu18
• rtx 3070
• cuda 11.1

因为要装的东西很多，我这里直接创建了一个conda虚拟环境，避免和我之前的环境产生冲突：

conda create -n yolox python=3.7    # 创建环境

source activate yolox             # 激活环境

然后按照官方给的安装示例，配置必须的包：

git clone git@github.com:Megvii-BaseDetection/YOLOX.git
cd YOLOX
pip3 install -U pip 
pip3 install -r requirements.txt              # 在requirements.txt 里我把torch注释掉了
python3 setup.py develop

我这里独自pip install安装的pytorch，版本是1.8，因为默认requirements.txt安装的话是安装最新版的1.9，而我在跑代码的时候因为1.9版本的问题报错了，张量cuda加载不上，如果你遇到同样的问题的话建议独自安装1.8版本的torch。

然后安装nvidia混合精度库apex：

git clone https://github.com/NVIDIA/apex
cd apex
pip3 install -v --disable-pip-version-check --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./

在安装apex的时候遇到了新问题，我的电脑cuda版本是11.1，而我安装的pytorch cuda版本是cuda 11.0 的，所以编译没通过，解决办法是安装cuda 11.1的torch版本，这两个版本要同步，torch离线包的网站有对应的whl文件，下载下来重新安装即可。如果这一步没报错的话，可以直接跳过这段话。

再然后安装pycocotools，我这一步比较顺利，直接通过：

pip3 install cython; pip3 install 'git+https://github.com/cocodataset/cocoapi.git#subdirectory=PythonAPI'

二、demo测试

环境配置完后，来运行一下demo测试，看看环境是否安装成功。在hub中将权重文件下载下来：

有点大，耐心等待。下载下来后我在YOLOX主目录中创建了一个weights文件夹，用于存放这些权重文件，然后直接运行：

python tools/demo.py image -n yolox-s -c /path/to/your/yolox_s.pth.tar --path assets/dog.jpg --conf 0.3 --nms 0.65 --tsize 640 --save_result               # 这是官方的示例

python tools/demo.py image -n yolox-s -c weights/yolox_s.pth.tar --path assets/bus.jpg --conf 0.3 --nms 0.65 --tsize 640 --save_result           #这是我的

其中：

• image是推理模式，如果是输入视频则为video
• -n 是模型的名字
• -c 为权重文件地址
• –path是测试的图片路径
• –conf 置信度阈值
• –nms nms的iou阈值
• –tsize 测试图片大小
• –save_result 是否保存推理结果

运行成功示例如下（注意：他这里的demo代码出错了也不会报错，直接跳出终止，所以如果你运行了发现没有如下的运行结果，可能是程序终止了，需要自己排查，我在这里遇到的问题是img的张量cuda加载不上，后面排查是因为torch1.9的问题，就换了1.8，正如前面所说）：

推理后的图片保存在data文件夹中，效果还不错的样子：

这步之后环境基本上没问题了，接下来准备数据开始训练吧。

三、数据准备和配置修改

官方配置教程

程序中提供了voc和coco两种格式的数据训练，我这里用的是voc格式训练，voc格式的文件结构如下：

├── data #手动创建data、VOCdevkit、VOC2007、Annotations、JPEGImages、ImageSets、Main这些文件夹
│   ├── VOCdevkit
│   │   ├── VOC2007
│   │   │   ├── Annotations #把test.txt、trainval.txt对应的xml文件放在这
│   │   │   ├── JPEGImages #把test.txt、trainval.txt对应的图片放在这
│   │   │   ├── ImageSets
│   │   │   │   ├── Main
│   │   │   │   │   ├── test.txt 
│   │   │   │   │   ├── trainval.txt

我的初始文件夹只有两个，通常标注完后只有这两个文件夹：

所以新建一个ImageSets文件，并在其子目录再新建一个Main文件，然后运行如下代码，划分数据集，得到test.txt和trainval.txt文件（路径看着修改）：

import os
import random

trainval_percent = 0.1
train_percent = 0.9
xmlfilepath = '/home/cai/data/VOCdevkit/VOC2007/Annotations'
txtsavepath = '/home/cai/data/VOCdevkit/VOC2007/ImageSets'
total_xml = os.listdir(xmlfilepath)

num = len(total_xml)
list = range(num)
tv = int(num * trainval_percent)
tr = int(tv * train_percent)
trainval = random.sample(list, tv)
train = random.sample(trainval, tr)

ftest = open('/home/cai/data/VOCdevkit/VOC2007/ImageSets/test.txt', 'w')
ftrain = open('/home/cai/data/VOCdevkit/VOC2007/ImageSets/trainval.txt', 'w')

for i in list:
    name = total_xml[i][:-4] + '\n'
    if i in trainval:
        ftest.write(name)
    else:
        ftrain.write(name)

ftrain.close()
ftest.close()

要严格按照这个格式制作数据集，因为程序中是按照这个结构读取，然后修改exps/example/yolox_voc/yolox_voc_s.py文件三个地方，把你的类别和文件路径改过来。



在我写这篇文章的时候，这个文件还有其他三个报错的地方，如下：



解决办法分别在yolox/data/datasets/init.py、yolox/evaluators/init.py两个文件中添加：


然后修改yolox/data/datasets/voc_classes.py，这里加上自己的类别，注释掉原始类别：

修改完过后直接运行如下命令开始训练：

python tools/train.py -f exps/example/yolox_voc/yolox_voc_s.py -d 1 -b 16 --fp16 -o -c weights/yolox_s.pth.tar

• -d 使用多少张显卡训练
• -b 批次大小
• –fp16 是否开启半精度训练

yolox_voc_s.py文件里的EXP类是继承于yolox_base.py，有一些参数可以看着修改，运行过后进入如下界面开始训练，打印了一堆的东西，我只截取了一部分：


然后你可能会遇到这样的问题，在迭代第二轮的时候迭代出来的内容为None，这时候需要去yolox/data/data_prefetcher.py文件下修改：

注释掉后面三行，改为pass。然后继续训练，正常训练后的输出如下，每两个轮次评估一次，可以在yolox_voc_s.py中修改self.eval_interval = 2来选择，这里只用了两百张左右的图片训练，可以看到收敛的速度非常的快，在两轮后已经有这么高的精度了：

紧接着，模型训练完后保存在data/YOLOX_outputs/yolox_voc_s文件夹中，只需要best模型即可，因为这个模型保存的是最好的精度：

四、模型测试

紧接着是测试，同样是运行之前的demo.py文件，但是需要修改一下demo.py，导入VOC_classes，然后修改可视化函数的传参

修改exps/default/yolox_s.py文件

修改你的类别数，不加的话会按照默认的coco 80类初始化网络，然后训练的权值无法加载会中断程序。这里我把best的权重放到了weights文件夹下，如果不修改运行后像这样：

不报错也不提示，坑爹啊

不过修改后就可以了，正确运行的界面如下。我们来看下我两百张图片训练出来的效果，用的还是之前筷子的数据集：

效果感觉也还行吧，和v5s的效果差不多，精度也比较接近，总的来说，yolox也对得起这个名字，毕竟人家不仅开源了训练代码，部署的也给你放出来了，卷啊卷啊卷。。。只不过这个代码，真是一堆的问题，差点给我整吐了。万幸还是跑出来了，大家遇到问题的话需要耐心调试，可以先去查yolox的issues，里面可能有你遇到的问题的解决方案，实在没有的话只能print。

五、最后补充

因为是跑通了之后才写的文章，有一些细节可能忘记加上去了，踩的坑有点多，遇到的问题也比较多，大部分都是比较容易解决的，我隐约感觉有一些修改的地方忘记写上去了，如果是用我提供的筷子数据集训练的话，可能会遇到这样的问题，因为我的有一些标签是通过程序生成的（之前的yolov5半标注），所以在xml文件中没有这两个属性：

上面是用labelimg标注的标准xml文件，下面是程序生成的，没有pose和truncated属性：

但是在yolox/evaluators/voc_eval.py中，解析xml的时候是有这两个属性的，我看了前后的代码，发现这两个属性可有可无，所以如果报错的话可以将其注释掉：

写的有点乱，但是大致的训练过程就是这样，中间大大小小的问题实在是太多了，记起来的都已经写出来了，可能会有遗漏。算法是好算法，不过这个代码在易用性上实在是有点不可恭维，希望后续作者能再优化好一些吧。