用Mask R-CNN训练自己的COCO数据集(Detectron)

用Mask R-CNN训练自己的COCO数据集(Detectron)

0、简介

今年年初,Facebook AI 研究院(FAIR)开源了 Detectron,业内最佳水平的目标检测平台。据介绍,该项目自 2016 年 7 月启动,构建于 Caffe2 之上,目前支持大量机器学习算法,其中包括 Mask R-CNN(何恺明的研究,ICCV 2017 最佳论文)和 Focal Loss for Dense Object Detection,(ICCV 2017 最佳学生论文)。


本文以 Airbus Ship Detection Challenge 为例,讲述如何基于Detectron内的Mask R-CNN训练自己的数据。


1、环境准备(Docker)

Detectron基于Caffe2,环境极为复杂,建议使用docker来部署。官方提供了镜像的dockerfile:facebookresearch/Detectron,编译即可生成本地镜像。我在此基础上做了一些修改,主要是安装了vim、添加了阿里软件源、安装pycocotools,大家可以通过下述命令获取我的镜像(缺点是这个镜像是静态的):

$ docker pull pascal1129/detectron:caffe2_cuda9_aliyun

之后生成容器,挂载目录什么的,按照你们的习惯来就行,示例如下:

$ nvidia-docker run -dit \
    --name pascal_ron \
    -v /home/pascal/data:/data \
    pascal1129/detectron:caffe2_cuda9_aliyun \
    bash
$ nvidia-docker exec -it pascal_ron bash

更多有关Docker的教程信息:

https://github.com/pascal1129/public_notes/blob/master/Docker.mdgithub.com
pascal1129/public_noteshttps://github.com/pascal1129/public_notes/blob/master/Docker.mdgithub.com


2、Docker环境测试

# 测试Caffe2环境是否配置正确,显卡是否可以检测到
python -c 'from caffe2.python import core' 2>/dev/null && echo "Success" || echo "Failure"
python -c 'from caffe2.python import workspace; print(workspace.NumCudaDevices())'
python /detectron/detectron/tests/test_batch_permutation_op.py

# 检测pycocotools是否正常
python -c 'from pycocotools.coco import COCO' 2>/dev/null && echo "Success" || echo "Failure"

使用预训练模型跑demo程序

python tools/infer_simple.py \
    --cfg configs/12_2017_baselines/e2e_mask_rcnn_R-101-FPN_2x.yaml \
    --output-dir /tmp/detectron-visualizations \
    --image-ext jpg \
    --wts https://s3-us-west-2.amazonaws.com/detectron/35861858/12_2017_baselines/e2e_mask_rcnn_R-101-FPN_2x.yaml.02_32_51.SgT4y1cO/output/train/coco_2014_train:coco_2014_valminusminival/generalized_rcnn/model_final.pkl \
    demo

3、制作自己的coco标准数据集

kaggle airbus challenge提供了一个28GB的训练集,使用我开源的代码可以制作COCO标准数据集,代码如下:pascal1129/airbus_ship_detection


4、配置数据集文件

添加数据集信息:

打开 /detectron/detectron/datasets/dataset_catalog.py 添加下述内容(供示例)

    'airbus_2018_train': {
        _IM_DIR:
            '/data/airbus_dataset/input/ships_train2018',
        _ANN_FN:
            '/data/airbus_dataset/input/annotations/instances_ships_train2018.json'
    },

添加类别信息:

打开/detectron/detectron/datasets/dataset_catalog.py 将种类更换为自己的类别


5、修改配置文件

我们假设使用/detectron/configs/12_2017_baselines/e2e_mask_rcnn_R-50-FPN_1x.yaml这个配置文件,主要是按需修改如下几个配置:

NUM_CLASSES: 2
NUM_GPUS: 2
训练集、测试集
DATASETS: ('airbus_2018_train',)

运行程序,开始训练

mkdir out_dir
python2 tools/train_net.py --cfg configs/12_2017_baselines/e2e_mask_rcnn_R-50-FPN_1x.yaml OUTPUT_DIR out_dir

如果中途报错,查阅资料修改即可,这里不多做解释了,当时我也查了不少资料,问题也比较复杂。比如说nan错误,可以参见这个:

retinanet loss go to NaN · Issue #81 · facebookresearch/Detectrongithub.com

训练时间可长可短,短则几个小时,长则几天

6、infer查看结果

使用tools/infer_simple.py运行即可

编辑于 2018-12-08

文章被以下专栏收录