训练自定义数据。
介绍如何使用 YOLOX 训练您自己的自定义数据。
我们举一个在 VOC 数据集上微调
YOLOX-S
模型的例子,以给出更清晰的指导。
0. 开始之前
克隆这个 repo 并按照 README 安装 YOLOX。
1. 创建自己的数据集
Step 1
首先准备您自己的带有图像和标签的数据集。
按照VOC的数据标注格式标注数据,数据目录如下:
VOCdevkit └── VOC2028 ├── Annotations ├── ImageSets │ └── Main └── JPEGImages
Step 2
将您的 VOC 数据集放在 $YOLOX_DIR/datasets
下 :
ln -s /path/to/your/VOCdevkit ./datasets/VOCdevkit
- 路径 “VOCdevkit” 将在下一节描述的 exp 文件中使用。具体来说,在
get_data_loader
和get_eval_loader
函数中。
2.创建你的Exp文件来控制一切
我们将模型中涉及的所有内容都放在一个单独的 Exp 文件中,包括模型设置、训练设置和测试设置。
完整的 Exp 文件位于 yolox/exp/yolox_base.py
。 可能每个exp都写太长,但是你可以继承基础exp文件,只覆盖改变的部分。
我们以 exps/example/yolox_voc/yolox_voc_s.py
文件为例。
我们这里选择了YOLOX-S
模型,所以我们应该改变网络的深度和宽度。准备的 VOC 只有 2 个类,因此我们也应该更改 num_classes
。
-
修改类别数量
在
init()
方法中更改def __init__(self): super(Exp, self).__init__() self.num_classes = 2 ##类别 self.depth = 0.33 self.width = 0.50 self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0]
-
修改训练集目录
在
get_data_loader()
方法中更改dataset = VOCDetection( data_dir=os.path.join(get_yolox_datadir(), "VOCdevkit"), image_sets=[('2028', 'trainval')], ## 训练集目录 img_size=self.input_size, preproc=TrainTransform( rgb_means=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225), max_labels=50, ), )
-
修改测试集目录
在
get_eval_loader()
方法中更改valdataset = VOCDetection( data_dir=os.path.join(get_yolox_datadir(), "VOCdevkit"), image_sets=[('2028', 'test')], ## 测试集目录 img_size=self.test_size, preproc=ValTransform( rgb_means=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225), ), )
-
修改类别
类别文件在
yolox/data/datasets/voc_classes.py
文件中,将自己对应的类别写入到这里即可。VOC_CLASSES = ( "helmet", "head" )
训练的其他设置可以参考
yolox/exp/yolox_base.py
的get_data_loader、get_eval_loader和get_evaluator了解更多详情。
3. 训练
除特殊情况外,我们始终建议使用我们的 COCO 预训练权重 来初始化模型。
获得我们提供的 Exp 文件和 COCO 预训练权重后,您可以通过以下命令训练自己的模型:
python tools/train.py -f /path/to/your/Exp/file -d 8 -b 64 --fp16 -o -c /path/to/the/pretrained/weights
或者以YOLOX-S
训练 VOC 为例:
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.tar
-
-d/--device
参数指定使用的的 GPU 数量 -
-b/--batch-size
参数推荐设置为-d
的 8 倍
不用担心预训练权重和你自己模型的检测头形状不同,我们会处理
注意:更改测试集的时候,需要删除 datasets/VOCdevkit 下的***annotations_cache***文件,该文件存储的时候测试集数据GT,如更改数据集,之前保存的有,则还是调用之前生成的文件,导致训练错误。
4. 获得最佳训练效果的提示
由于 YOLOX 是一个只有几个超参数的无锚检测器,大多数情况下可以在不改变模型或训练设置的情况下获得良好的结果。
因此,我们始终建议您首先使用所有默认训练设置进行训练。
如果一开始你没有得到好的结果,你可以考虑采取一些步骤来改进模型。
模型选择 我们为移动部署提供YOLOX-Nano
、YOLOX-Tiny
和YOLOX-S
,而YOLOX-M
/L
/X
用于云或高性能 GPU 部署。
如果您的部署遇到任何兼容性问题。我们推荐 YOLOX-DarkNet53
.
训练配置 如果您的训练过早过拟合,那么您可以在您的 Exp 文件中减少 max_epochs 或减少 base_lr 和 min_lr_ratio:
# -------------- training config --------------------- #
self.warmup_epochs = 5
self.max_epoch = 300
self.warmup_lr = 0
self.basic_lr_per_img = 0.01 / 64.0
self.scheduler = "yoloxwarmcos"
self.no_aug_epochs = 15
self.min_lr_ratio = 0.05
self.ema = True
self.weight_decay = 5e-4
self.momentum = 0.9
Aug 配置 您还可以更改增强的程度。
一般来说,对于小模型,你应该弱化aug,而对于大模型或小数据集,你可以在exp文件中增强aug:
# --------------- transform config ----------------- #
self.degrees = 10.0
self.translate = 0.1
self.scale = (0.1, 2)
self.mscale = (0.8, 1.6)
self.shear = 2.0
self.perspective = 0.0
self.enable_mixup = True
设计你自己的检测器 您可以参考我们的 Arxiv 论文,了解有关设计您自己的检测器的详细信息和建议。