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

Posted yinqinggong

tags:

篇首语:本文由小常识网(cha138.com)小编为大家整理,主要介绍了YOLOX推理系列4-使用YOLOX训练自己的数据集相关的知识,希望对你有一定的参考价值。


前言

前面的文章,笔者在系统环境下使用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 git@github.com: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推理系列打完收工

以上是关于YOLOX推理系列4-使用YOLOX训练自己的数据集的主要内容,如果未能解决你的问题,请参考以下文章

yolox_nano训练和NCNN安卓部署总结(自己的数据集)

[深度学习][原创]yolox快速训练自己数据集方法

YOLOX训练自己的数据集(超详细)

yolox训练自己的数据

yolox训练自己的数据

yolox训练自己的数据