使用 Mask-RCNN Benchmark 训练自定义数据集

介绍

  Facebook AI Research 开源了 Faster R-CNN 和 Mask R-CNN 的 PyTorch 1.0 实现框架：MaskRCNN-Benchmark。他的性能相对更快，并且占用更少的 GPU 内存，它的亮点如下：

• PyTorch 1.0: 相当或者超越 Detectron 准确率的 RPN、Faster R-CNN 和 Mask R-CNN 的实现；
• 非常快: 训练速度是 Detectron 的两倍，是 mmdection 的 1.3 倍。
• 节省内存: 在训练过程中使用的 GPU 内存比 mmdetection 少大约 500MB；
• 使用多 GPU 训练和推理；
• 批量化推理:可以在每 GPU 每批量上使用多张图像进行推理；
• 支持 CPU inference: 可以在 CPU 上运行。
• 提供几乎所有参考 Mask R-CNN 和 Faster R-CNN 配置的预训练模型，具有 1x 的 schedule。

  这篇文章主要是记录我使用 Mask R-CNN benchmark 框架训练自定义数据集的过程，总的来说还是比较容易上手的，当然也有一些问题出现。下面便是使用的情况。

安装

我的开发环境如下:

• 系统: Ubuntu 18.04LTS
• Python 环境: miniconda3
  • conda: 4.7.12
  • Python: 3.7.4
• GCC: 7.3.0
• CUDA: 10.0
• pip: 19.3.1

项目的要求如下:

• PyTorch 1.0 from a nightly release. It will not work with 1.0 nor 1.0.1. Installation instructions can be found in https://pytorch.org/get-started/locally/
• torchvision from master
• cocoapi
• yacs
• matplotlib
• GCC >= 4.9
• OpenCV
• CUDA >= 9.0

接下来跟着执行安装命令:

# create the development environment
conda create --name detection
source activate detection
# install the dependences
conda install ninja yacs cython matplotlib tqdm opencv-python pycocotools
conda install -c pytorch pytorch-nightly torchvision cudatoolkit=10.0

---

# install the core projects
export INSTALL_DIR = $PWD
cd $INSTALL_DIR 
# install apex
git clone https://github.com/NVIDIA/apex.git
cd apex
python setup.py install --cuda_ext --cpp_ext
cd $INSTALL_DIR
# install PyTorch Detection 
git clone https://github.com/facebookresearch/maskrcnn-benchmark.git
cd maskrcnn-benchmark
python setup.py build develop
unset INSTALL_DIR

准备数据集

  安装完成后，下一步是要准备训练数据和验证数据。
  maskrcnn-benchmark默认是使用COCO数据集的，因此需要把自己的数据集改成COCO的Detection格式。该数据集的标准是使用JSON文件存储，具体内容可以参考COCO数据集的标注格式一文。
  COCO数据集的Detection部分主要的JSON结构体类型如下所示：

annotations.json

{
    "root":{
        "images": [
            "0": {
                "id": int, //每张图片都有唯一的id标识
                "file_name": str, //图片文件路径
                "width": int,
                "height": int
            },
        ],
        "categories": [
            "0": {
                "id": int,
                "name": str,
                "count": int  //每种category的数量
            },
        ],
        "annotations": [
            "0": {
                "id": int, //每个annotation都有唯一的id标识
                "image_id": int, //对应于图片的id
                "segmentation": RLE or [polygon],  //用多边形表示mask区域
                "category_id": int,  //对应的category id
                "iscrowd": 0 or 1,  //表示是否是一组对象，如果为1，则segmentation用RLE格式，如果为0，则用polygons格式
                "area": float,  //区域面积
                "bbox": [x,y,width,height]  //通过这四个坐标定位bounding box
            },
        ]
    }
}

  参考上面的格式分别生成训练集和验证集的json标注文件，随后将数据集文件软链接到dataset文件夹目录下。目录格式可以参考下面:

├── datasets
│   └── coco #可以任意修改名称
│       ├── annotations
│  		│  ├── instances_train2014.json #训练集标注文件
│  		│  └── instances_val2014.json #验证集标注文件
│       ├── train2014  #存放训练集图片
│       └── val2014  #存放验证集图片

修改配置文件

  在准备好数据文件之后，便可以通过修改配置文件来开始训练数据集，这里涉及的配置文件如下：

• 模型参数配置文件：configs/e2e_mask_rcnn_R_101_FPN_1x.yaml
• 数据集路径配置文件： maskrcnn_benchmark/config/path_catalog.py

数据路径文件

  需要修改的地方如下：

path_catalog.py

# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
"""Centralized catalog of paths."""

import os
from copy import deepcopy

class DatasetCatalog(object):
    DATA_DIR = "datasets"
    DATASETS = {
        "coco_2017_train": {
            "img_dir": "coco/train2017",
            "ann_file": "coco/annotations/instances_train2017.json"
        },
        "coco_2017_val": {
            "img_dir": "coco/val2017",
            "ann_file": "coco/annotations/instances_val2017.json"
        },
        ···
    }

然后在DETASETS里面按照同样的格式修改image文件夹和annotation文件的位置，需要与自己的数据集相互对应。修改好文件位置之后，则修改模型参数文件。

模型参数文件

  模型参数文件里面的内容主要是网络结构、数据集配置和训练方式等，由于官方的模型是在8个GPU的环境下训练的，因此我们需要调小训练的batch_size，防止显存溢出。同时调整学习率，使网络能够正常训练。我的配置文件如下所示:

1

启动训练

  进入maskrcnn-benchmark根目录下，激活虚拟环境，并执行训练脚本。

cd maskrcnn-benchmark
source activate detection
python tools/train_net.py --config-file configs/adas_e2e_mask_rcnn_R_101_FPN_1x.yaml

每隔一定的iteration会打印loss信息，用于判断训练的结果是否走偏，从而节省时间，这里我使用的是默认值20。最后就是验证结果，该步骤会在训练完成后自动运行，我们看一下结果：