使用mmdetection训练自己的数据集(记录)
简介
商汤科技(2018 COCO目标检测挑战赛冠军)开源了基于Pytorch的目标检测工具mmdetection,支持Mask-RCNN,Fast-RCNN,Faster-RCNN等主流的目标检测框架。
作者称mmdetection有三点优势:测试结果、训练速度、所需显存都有所改善。
下面将记录一下我的测试过程。
准备自己的数据集
mmdetection支持coco格式和voc格式的数据集,分别介绍这两种数据集的使用方式。
coco数据集
官方推荐coco数据集按照以下目录形式存储,以coco2017数据集为例
1 | mmdetection |
推荐以软连接的方式创建data文件夹,下面是创建软连接的步骤
1 | cd mmdetection |
其中,$COCO_ROOT需改为你的coco数据集根目录
voc数据集
与coco数据集类似,将voc数据集按照以下的目录形式存储,以VOC2007为例
1 | mmdetection |
同样推荐以软连接的方式创建
1 | cd mmdetection |
其中,$VOC2007_ROOT需改为你的VOC2007数据集根目录
数据集的处理过程
切图->json文件分类->img放入相应json文件中->按照各类数据集数量均衡合成数据集(通过各类数量可视化)->通过labelme过一遍数据集有无错误->将img和json分train和val->生成相应的mask->生成一个json文件,再修正label中train和test的类别对应格式。文件目录如上图所示,准备进行训练。
修改相关文件参数
mmdet /core/evaluation/class_name.py中的coco_classes
mmdet/datasets/coco.py中的class CocoDataset(CustomDataset)中的CLASSES
mmdet/datasets/coco.py中的def evaluate()中的 classwise=true
train出来—>pyconfig文件参数修改,如:image_scale,epoch,lr,batch-size等
开始训练
官方推荐使用分布式的训练方式,这样速度更快,如果是coco训练集,修改CONFIG_FILE中的pretrained参数,改为你的模型路径,然后运行下方代码
1 | bash ./tools/dist_train.sh <CONFIG_FILE> <GPU_NUM> [optional arguments] |
如果不想采用分布式的训练方式,或者你只有一块显卡,则运行下方的代码
1 | python tools/train.py <CONFIG_FILE> --gpus <GPU_NUM> --work-dir <WORK_DIR> |
如果一切顺利的话,你的模型应该就开始训练了
准备测试
使用如下代码进行测试
1 | bash tools/dist_test.sh <CONFIG_FILE> path/to/model <GPU_NUM> —format-only —eval-options "jsonfile_prefix=./xxx" |
生成.seg和.bbox的json文件,并进行可视化测试或者通过其他测试方式进行具体测试分析。