使用YoloX进行自定义训练

昨天因为要准备比赛,所以搞了一阵yolox的custom_train,效果很ok,遂分享出来。

下载并配置YoloX

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

怎么clone就不需要我多言了。

pip安装完环境,我们需要clone一下NVIDIA的apex,最后执行

1
python3 setup.py develop

基本的结构介绍

介绍一下我们自定义训练主要用到的几个文件:

模型主要部分:./yolox

各种工具:./tools

其中我们需要使用train.py来训练,demo.py来预测。

命令分别为:

1
python tools/train.py -f exps/example/yolox_voc/yolox_voc_s.py -d 8 -b 64 --fp16 -o -c /path/to/yolox_s.pth [--cache]
1
python tools/demo.py image -f exps/default/yolox_s.py -c /path/to/your/yolox_s.pth --path assets/dog.jpg --conf 0.25 --nms 0.45 --tsize 640 --save_result --device [cpu/gpu]

或者视频:

1
python tools/demo.py video -n yolox-s -c /path/to/your/yolox_s.pth --path /path/to/your/video --conf 0.25 --nms 0.45 --tsize 640 --save_result --device [cpu/gpu]

注意:之后要把./yolox/exp/yolox_base.py里面的num_classesdataset改掉。

注意:如果之后模型报不能找到”yolox“这个module的话,需要在前面添加

1
2
import sys

sys.path.append(r'你的/根/目录')

使用VOC类型的数据集训练yolo_s模型需要用到的Exp:./exps/example/yolox_voc/yolox_voc_s.py

这个文件主要定义了我们训练的Exp,我们自定义训练需要修改的主要包括分类数(num_classes)、数据集(dataset)。

数据集:./datasets

我们需要在里面放置VOC或者COCO格式的数据集。

我是用的是自制VOC结构的数据集,结构如下:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
 .
└── VOCdevkit
└── VOC2007
├── Annotations
│ ├── 003002_0.xml
│ ├── 003002_1.xml
│ ├── 003008_1.xml
│ └── .......
├── ImageSets
│ └── Main
│ ├── test.txt
│ ├── train.txt
│ ├── trainval.txt
│ └── val.txt
└── JPEGImages
├── 003002_0.jpg
├── 003002_1.jpg
├── 003008_1.jpg
└──.......

而我初始的数据集只有jpg和xml,所以需要转换。

转换数据集:

先把原始图片目录改成这样,创建用于生成数据集的xml2voc2007.py

1
2
3
4
5
6
7
8
9
10
. 
└── data
├── 003002_0.jpg
├── 003002_0.xml
├── 003002_1.jpg
├── 003002_1.xml
├── 003008_1.jpg
├── 003008_1.xml
└── .......
└── xml2voc2007.py

附上用xml和jpg生成VOC格式数据集的代码:

xml2voc2007.py >folded
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90

#命令行执行: python xml2voc2007.py --input_dir data --output_dir VOCdevkit

import argparse
import glob
import os
import random
import os.path as osp
import sys
import shutil

percent_train = 0.9

# 主程序执行

def main():
parser = argparse.ArgumentParser(
formatter_class=argparse.ArgumentDefaultsHelpFormatter
)
parser.add_argument("--input_dir", default="data", help="input annotated directory")
parser.add_argument("--output_dir", default="VOCdevkit", help="output dataset directory")
args = parser.parse_args()

if osp.exists(args.output_dir):
print("Output directory already exists:", args.output_dir)
sys.exit(1)
os.makedirs(args.output_dir)
print("| Creating dataset dir:", osp.join(args.output_dir, "VOC2007"))

# 创建保存的文件夹
if not os.path.exists(osp.join(args.output_dir, "VOC2007", "Annotations")):
os.makedirs(osp.join(args.output_dir, "VOC2007", "Annotations"))
if not os.path.exists(osp.join(args.output_dir, "VOC2007", "ImageSets")):
os.makedirs(osp.join(args.output_dir, "VOC2007", "ImageSets"))
if not os.path.exists(osp.join(args.output_dir, "VOC2007", "ImageSets", "Main")):
os.makedirs(osp.join(args.output_dir, "VOC2007", "ImageSets", "Main"))
if not os.path.exists(osp.join(args.output_dir, "VOC2007", "JPEGImages")):
os.makedirs(osp.join(args.output_dir, "VOC2007", "JPEGImages"))

# 获取目录下所有的.jpg文件列表
total_img = glob.glob(osp.join(args.input_dir, "*.jpg"))
print('| Image number: ', len(total_img))

# 获取目录下所有的joson文件列表
total_xml = glob.glob(osp.join(args.input_dir, "*.xml"))
print('| Xml number: ', len(total_xml))

num_total = len(total_xml)
data_list = range(num_total)

num_tr = int(num_total * percent_train)
num_train = random.sample(data_list, num_tr)

print('| Train number: ', num_tr)
print('| Val number: ', num_total - num_tr)

file_train = open(
osp.join(args.output_dir, "VOC2007", "ImageSets", "Main", "train.txt"), 'w')
file_val = open(
osp.join(args.output_dir, "VOC2007", "ImageSets", "Main", "val.txt"), 'w')

for i in data_list:
name = total_xml[i][:-4] + '\n'
if i in num_train:
file_train.write(name[5:])
else:
file_val.write(name[5:])

file_train.close()
file_val.close()

if os.path.exists(args.input_dir):
# root 所指的是当前正在遍历的这个文件夹的本身的地址
# dirs 是一个 list,内容是该文件夹中所有的目录的名字(不包括子目录)
# files 同样是 list, 内容是该文件夹中所有的文件(不包括子目录)
for root, dirs, files in os.walk(args.input_dir):
for file in files:
src_file = osp.join(root, file)
if src_file.endswith(".jpg"):
shutil.copy(src_file, osp.join(args.output_dir, "VOC2007", "JPEGImages"))
else:
shutil.copy(src_file, osp.join(args.output_dir, "VOC2007", "Annotations"))

print('| Done!')

if __name__ == "__main__":
print("—" * 50)
main()
print("—" * 50)

开始训练自己的数据集

以VOC格式的数据集举例:

先在./yolox/data/datasets/voc_classes.py里,把VOC_CLASSES改成你的数据集类别

接着把./exps/example/yolox_voc/yolox_voc_s.py里面的num_classes改成你的类别数

好了,然后我们新建个文件夹,去仓库里把预训练权重下载下来,准备工作就完事了。

训练命令:

1
python3 tools/train.py -f exps/example/yolox_voc/yolox_voc_s.py -d 1 -b 8 --fp16 -o -c models/yolox_s.pth`

其中:

-d是num_workers,

-b是batch_size,

-c是cache,就是之前下载的权重。

好的,这样我们就可以开始训练了。

可视化训练过程

yolox自带有tensorboard输出,我们可以在安完tensorboard以后,进入./YOLOX_outputs/yolox_voc_s/tensorboard里面,输入

1
tensorboard --logdir ./ --bind_all

来可视化训练过程。

评估效果

参照上面demo的命令。

Author

holk

Posted on

2022-04-06

Updated on

2024-02-17

Licensed under

Comments