Article / 文章中心

YOLOX-0.1.0自定义训练+CPU修改方式

发布时间:2022-10-11 点击数:1656

1、主训练模块train.py在tools文件夹下(为什么不放在主目录?),训练命令为python tools/train.py -f exps/example/yolox_voc/yolox_s.py

YOLOX/tools/

 

2、train.py 中必须加入import sys sys.path.append(r'D:\YoloX')才能正常运行,即主目录的路径

train.py

 

3、模型深度控制模块yolox_s.py在exps/example/yolox_voc文件夹下,训练命令后半部分这么长就是这么来的,可以复制原模块再自定义内容

yolox_voc_s.py为原模块

 

4、数据集导入模块voc.py在yolox/data/datasets文件夹下,如果复制一份再重命名并自定义,必须修改datasets文件夹中的__init__.py的from .voc import VOCDetection为自己的文件名

YOLOX/yolox/data/datasets
YOLOX/yolox/data/datasets/__init__.py

 

5、数据集导入模块voc.py的get_data_loader()为导入训练数据,get_eval_loader()为导入验证数据,如果只是进行一些代码测试,不想要测试集,那么可以将测试集路径设置为与训练集路径一致

get_data_loader()
get_eval_loader()

 

6、类别设置模块voc_classes.py在yolox/data/datasets文件夹下,如果复制一份再重命名并自定义,必须修改voc.py中的from .voc_classes import VOC_CLASSES中的voc_classes为自己的文件名

YOLOX/yolox/data/datasets/voc.py

 

7、数据集必须为VOC格式,放在YOLOX/datasets/下,文件夹名称为VOCdevkit/VOC2007/,其中三个文件夹Annotations存放xml标签,ImageSets/Main存放train.txt,即图片名称,不能带后缀,JPEGImages存放图片

 

8、如果想在本地用CPU进行测试,必须安装CUDA版torch,否则会报错AssertionError: Torch not compiled with CUDA enabled

然后在tools/train.py中assert num_gpu <= get_num_devices()后插入num_gpu = 0

再在YOLOX/yolox/core/trainer.py的Trainer类的__init__函数中把self.device = "cuda:{}".format(self.local_rank)注释掉,插入self.device = "cpu"

再在YOLOX/yolox/core/trainer.py中的train_one_iter函数中的targets = targets.to(self.data_type)后插入inps = inps.to(self.device)与targets = targets.to(self.device)

并且把下面的with torch.cuda.amp.autocast(enabled=self.amp_training):注释掉

这样就可以在本地用CPU进行一些测试了

tools/train.py
YOLOX/yolox/core/trainer.py/__init__()
YOLOX/yolox/core/trainer.py/train_one_iter()

 

总体来说代码对于自定义训练与测试很不友好,没有Ultralytics的YOLOv5简单方便易懂