YOLOX推理系列4-使用YOLOX训练自己的数据集
使用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 [email protected]: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推理系列打完收工