[Python知识库]Yolox训练自己的数据集教程

本文环境：ubuntu20.04+cuda11.0+pytorch1.7.1

?https://github.com/Megvii-BaseDetection/YOLOX

1. 配置环境

（1）创建虚拟环境并安装代码所需的库文件

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

??????? source activate yolox???????????? # 激活环境

??????? cd YOLOX-main

????????pip3 install -r requirements.txt??? #把torch注释掉了

（2）通过setup.py安装一些库文件

????????输入命令：python3 setup.py develop

（3）安装apex文件

????????git clone https://github.com/NVIDIA/apex
????????cd apex
????????python3 setup.py install

（4）下载权重文件

2.准备VOC数据集

放入yolo-main/data/YOLOX-main/datasets/VOCdevkit/VOC2007文件夹

3. 修改相应代码

（1）修改类别和类别数

??????? a. 将yolox/data/datasets/voc_classes.py中的标签信息，进行修改

????????注意：类别后面都要加逗号，例如“car”后面加了一个逗号“，”。

??????? b. 修改exps/example/yolox_voc/yolox_voc_s.py中的self.num_classes

????????c. 修改yolox/exp/yolox_base.py中的self.num_classes

（2）修改训练集信息

????????a. 修改exps/example/yolox_voc/yolox_voc_s.py中的VOCDection。

????????data_dir="/home/XX/YOLOX-main/datasets/VOCdevkit",

????????image_sets=[('2007', 'train')],

????????max_labels=100,

????????b. 修改exps/example/yolox_voc/yolox_voc_s.py中的get_eval_loader函数。

?????????valdataset = VOCDetection(
??????????? data_dir="/home/XX/YOLOX-main/datasets/VOCdevkit",
??????????? image_sets=[('2007', 'val')],
??????????? img_size=self.test_size,
??????????? preproc=ValTransform(legacy=legacy),
??????? )

（3）修改train.py

????????修改epoch数量，学习率

4. 开始训练

将下载好的yolox_s.pth.tar放到YOLOX文件夹中，打开终端，在终端中输入：

python tools/train.py -f exps/example/yolox_voc/yolox_voc_s.py -d 0 -b 4 --fp16? -c yolox_s.pth

d表示gpu编号，b表示batch size，fp16是否开启半精度训练

5. 测试

利用eval.py进行测试

1）修改eval.py文件中的batch_size大小，可根据个人的显卡大小修改

2）修改yolox的配置文件，与训练时的配置文件一致

????????default=“/home/XX/YOLOX-main/exps/example/yolox_voc/yolox_voc_s.py”

3）修改测试时使用的模型，得分阈值、nms阈值等参数

??????? parser.add_argument("--conf", default=0.1, type=float, help="test conf")
??????? parser.add_argument("--nms", default=0.5, type=float, help="test nms threshold")
?? ? ?? parser.add_argument("--tsize", default=640, type=int, help="test img size")

4)修改测试集

把exps/example/yolox_voc/yolox_voc_s.py中的valdataset = VOCDetection(
??????????? data_dir="/home/ss/shanshan/seedposition-detection/YOLOX-main/datasets/VOCdevkit",
??????????? image_sets=[('2007', 'test')],#val改为test
??????????? img_size=self.test_size,
??????????? preproc=ValTransform(legacy=legacy),
??????? )

5）测试命令：计算AP

python tools/eval.py -f exps/example/yolox_voc/yolox_voc_s.py -d 0 -b 4 -c YOLOX_outputs/yolox_voc_s/best_ckpt.pth