YOLOX推理系列4-使用YOLOX训练自己的数据集

Posted 2022-01-01 yinqinggong

使用YOLOX训练自己的数据集

• 前言
• 一、安装Anaconda
• 二、安装YOLOX
• • 1.笔者掉入的坑
  • 2.开始安装
  • 3.测试YOLOX
• 三、训练COCO128
• 四、训练自己的数据集

---

前言

前面的文章，笔者在系统环境下使用OpenCV对YOLOX模型进行推理；在本节，我们需要安装虚拟环境使用YOLOX训练自己的数据集。

---

一、安装Anaconda

去Anaconda官网下载最新的版本

如果查询到版本号，也可使用命令行：

wget https://repo.anaconda.com/archive/Anaconda3-2021.11-Linux-x86_64.sh

然后执行安装：

sudo bash Anaconda3-2021.11-Linux-x86_64.sh

安装完成后验证

如果没有出现则需要添加环境变量：

sudo vim ~/.bashrc

在最后添加：

export PATH=$PATH:/home/zrd/anaconda3/bin

路径需要换成自己的，然后执行

source ~/.bashrc

二、安装YOLOX

1.笔者掉入的坑

在YOLOX的系统推理环境中，笔者安装了最新的CUDA11.5和对应的cuDNN版本，顺利的使用了GPU推理。所以在conda虚拟环境中，笔者也准备直接使用系统的环境的CUDA，然而事与愿违，无法调用CUDA进行GPU训练。

查询后发现目前pytorch的最高版本是1.10，而它能支持的CUDA最高版本是11.3，详见Pytorch官网，这样的话我们就只能在conda环境中重新安装CUDA10.2或者CUDA11.3。所以安装pytorch训练环境的准则是，pytorch和CUDA以及cuDNN版本需要匹配，感谢好友陈总对此的指导。

2.开始安装

一，创建&进入虚拟环境

conda create -n yolox python=3.8
conda activate yolox

二，安装pytorch1.10+cuda11.3(如果系统环境支持则不需要此步骤)

conda install pytorch torchvision torchaudio cudatoolkit=11.3 -c pytorch

三，按照官网安装YOLOX
YOLOX-github
Step1. Install YOLOX.

git clone [e-mail]:Megvii-BaseDetection/YOLOX.git
cd YOLOX
pip3 install -U pip && pip3 install -r requirements.txt
pip3 install -v -e .  # or  python3 setup.py develop

Step2. Install pycocotools.

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

3.测试YOLOX

下载yolox_s.pth权重文件，放到datasets目录下：

使用cpu推理：

python tools/demo.py image -n yolox-s -c datasets/yolox_s.pth --path assets/dog.jpg --conf 0.25 --nms 0.45 --tsize 640 --save_result --device cpu

使用gpu推理：

python tools/demo.py image -n yolox-s -c datasets/yolox_s.pth --path assets/dog.jpg --conf 0.25 --nms 0.45 --tsize 640 --save_result --device gpu

结果都保存在YOLOX_outputs目录下：

三、训练COCO128

提示：在训练自己的数据集之前可以先使用官方的小数据集比如coco128测试一下，看看是否环境有错，如果coco128可以训练的话再更换为自己的数据集。这样如果发生错误能更好的定位。

关于coco128：
coco128是coco数据集的前128张图片，一共有71个类而不是80，主要用来测试。
在官方Train Custom Data中给出了coco128的下载地址，和训练方法：

先不要修改任何文件，只把coco128下载到datasets目录即可：

使用exps/example/custom/yolox_s.py脚本训练，这个脚本就是为coco128数据集准备的，后面训练自己的数据集时可以据此修改：

训练脚本：

python tools/train.py -f exps/example/custom/yolox_s.py -d 1 -b 8 --fp16 -o -c /path/to/yolox_s.pth

-d 1：number of gpu devices
-b 8：batch size, the recommended number for -b is num-gpu * 8
训练中：

训练完成：

训练完成后会将权重保存在YOLOX_outputs/yolox_s/目录下，选择best_ckpt.pth进行推理测试：

python tools/demo.py image -f exps/example/custom/yolox_s.py -c YOLOX_outputs/yolox_s/best_ckpt.pth --path assets/dog.jpg --conf 0.25 --nms 0.45 --tsize 640 --save_result --device gpu

去YOLOX_outputs/yolox_s/vis_res目录下查看推理图片，发现效果很差，没有关系，因为coco128数据集太小了。

四、训练自己的数据集

建议仿造coco128做自己的数据集，包括目录结构（annotations train2017 val2017）和json文件名（instances_train2017.json instances_val2017.json）等，这样就不用修改配置文件的相关内容。
把自己的数据集同样放到datasets目录下，然后修改两处配置。
一，yolox/data/datasets/coco_classes.py，把类名修改成自己的。

二，exps/example/custom/yolox_s_zrd.py，仿造yolox_s.py建立一个自己的训练文件

开始训练：

python tools/train.py -f exps/example/custom/yolox_s_zrd.py -d 1 -b 8 --fp16 -o -c datasets/yolox_s.pth

训练中：

训练完成：

使用自己训练的模型推理：

python tools/demo.py image -f exps/example/custom/yolox_s_zrd.py -c YOLOX_outputs/yolox_s_zrd/best_ckpt.pth --path datasets/zrodo_coco/val2017/0020.jpg --conf 0.25 --nms 0.45 --tsize 640 --save_result --device gpu

不戴口罩还抽烟的强哥

---

至此，YOLOX推理系列打完收工