前一段时间Facebook开源了其用于目标检测、识别和分割的框架——Detectron，该框架携带了这一领域里的经典和最新算法，诸如Faster-RCNN、RetinaNet、Mask-RCNN，等等。这绝对是造福这一领域从业者和研究者的事情，当然，也成了这一领域绕不开的东西了。 Detectron自带了可以基于Pascal VOC以及COCO数据集的模型训练方法。不过Detectron十分易用，即使训练自己的数据集，操作起来也十分简单方便。本文对此进行一下简单介绍和总结。

1.安装Detectron以及Caffe2

Detectron是基于Caffe2开发的，因此，首先需要安装Caffe2，然后从GitHub上clone下来Detectron进行安装，Detectron的安装说明文档十分详细且易于操作，直接参考进行安装即可。

1.1 Caffe2

安装完成Caffe2后，在Terminal里运行下面代码，检查是否安装成功：

# To check if Caffe2 build was successful
python -c ‘from caffe2.python import core’ 2>/dev/null && echo “Success” || echo “Failure”

# To check if Caffe2 GPU build was successful
# This must print a number > 0 in order to use Detectron
python -c ‘from caffe2.python import workspace; print(workspace.NumCudaDevices())’

1.2 安装一些依赖及可能用到的COCO API

1.3 安装Detectron

从GitHub上将其clone下来，编译：

cd $DETECTRON/lib && make

然后检查是否安装成功：

python $DETECTRON/tests/test_spatial_narrow_as_op.py

2.配置自己的数据集

因为自定义自己的数据集比较繁琐，所以推荐使用Detectron自带的COCO数据集名称，并将自己的数据集转化为COCO数据集的格式，具体配置如下。

2.1 放置数据

将数据放置到$DETECTRON/lib/datasets/data下面，如animals数据集，并在其中以如下结构放置数据集的数据： ├── annos ├── annotations ├── classes.txt └── images 其中，annos放置你的数据集的原始的标注文件，可能是txt，或者csv格式；classes.txt放置你标注的类别名称，每行一个类别，不含背景；images放置数据集的原始图像文件。annotations预备放置与COCO数据集格式的标注文件。下面需要将我们自己的标注文件生成COCO格式的标注文件。

2.2 转化标注格式

因为我们自己的数据集的标注文件是txt格式，每行一个标注，形似如下：

老虎 25 36 128 159
猫 25 36 128 159
狮子 25 36 128 159

所以，我们需要将这种格式的标注转换为COCO格式。程序如下：

#!/usr/bin/python
# -- coding: utf-8 --
import json
import os
import sys

import cv2

if len(sys.argv) < 3:
print “Usage: python convert_to_detectron_json.py root_path phase split”
print “For example: python convert_to_detectron_json.py data train 100200”
exit(1)

root_path = sys.argv[1]
phase = sys.argv[2]
split = int(sys.argv[3])

dataset = {
‘licenses’: [],
‘info’: {},
‘categories’: [],
‘images’: [],
‘annotations’: []
}

with open(os.path.join(root_path, ‘classes.txt’)) as f:
classes = f.read().strip().split()

for i, cls in enumerate(classes, 1):
dataset[‘categories’].append({
‘id’: i,
‘name’: cls,
‘supercategory’: ‘beverage’
})

def get_category_id(cls):
for category in dataset[‘categories’]:
if category[‘name’] == cls:
return category[‘id’]

_indexes = sorted([f.split(‘.’)[0] for f in os.listdir(os.path.join(root_path, ‘annos’))])
if phase == ‘train’:
indexes = [line for line in _indexes if int(line) > split]
else:
indexes = [line for line in _indexes if int(line) <= split]

j = 1
for index in indexes:
im = cv2.imread(os.path.join(root_path, ‘images/‘) + index + ‘.jpg’)
height, width, _ = im.shape
dataset[‘images’].append({
‘coco_url’: ‘’,
‘date_captured’: ‘’,
‘file_name’: index + ‘.jpg’,
‘flickr_url’: ‘’,
‘id’: int(index),
‘license’: 0,
‘width’: width,
‘height’: height
})

anno_file = os.path.join(root_path, ‘annos/‘) + index + ‘.txt’
with open(anno_file) as f:
lines = [line for line in f.readlines() if line.strip()]

for i, line in enumerate(lines):
  parts = line.strip().split()
  cls = parts\[0\]
  x1 = int(parts\[1\])
  y1 = int(parts\[2\])
  x2 = int(parts\[3\])
  y2 = int(parts\[4\])
  width = max(0, x2 - x1)
  height = max(0, y2 - y1)
  dataset\['annotations'\].append({
      'area': width * height,
      'bbox': \[x1, y1, width, height\],
      'category\_id': get\_category_id(cls),
      'id': j,
      'image_id': int(index),
      'iscrowd': 0,
      'segmentation': \[\]
  })
  j += 1

folder = os.path.join(root_path, ‘annotations’)
if not os.path.exists(folder):
os.makedirs(folder)
json_name = os.path.join(root_path, ‘annotations/{}.json’.format(phase))
with open(json_name, ‘w’) as f:
json.dump(dataset, f)

可以点击这里去GitHub查看源代码。你可以根据自己数据集的格式，改写程序，将标注转化为COCO格式。

2.3 设置指向自己的数据集

修改$DETECTRON/lib/datasets/dataset_catalog.py文件，将其中描述COCO数据集的两个部分修改为如下所示：

‘coco_2014_train’: {
IM_DIR:
_DATA_DIR + ‘/animals/images’,
ANN_FN:
_DATA_DIR + ‘/animals/annotations/train.json’
},
‘coco_2014_val’: {
IM_DIR:
_DATA_DIR + ‘/animals/images’,
ANN_FN:
_DATA_DIR + ‘/animals/annotations/val.json’
},

其中，_DATA_DIR是Detectron的数据的根目录，即$DETECTRON/lib/datasets/data，其余顾名思义。

3.新建配置文件，开始训练

到这里，设置好了相应的数据集，下面需要指定训练模型使用该数据集，Detectron是通过配置文件来指定的。

3.1 新建配置文件

可以直接复制Detectron自带的配置文件，并稍作修改即可，在$DETECTRON根目录下打开Terminal复制：

cd configs
mkdir animals
cp getting_started/tutorial_1gpu_e2e_faster_rcnn_R-50-FPN.yaml animals/animals_1gpu_e2e_faster_rcnn_R-50-FPN.yaml

然后打开animals_1gpu_e2e_faster_rcnn_R-50-FPN.yaml进行编辑：

MODEL:
TYPE: generalized_rcnn
CONV_BODY: FPN.add_fpn_ResNet50_conv5_body
NUM_CLASSES: 5
FASTER_RCNN: True
NUM_GPUS: 1
SOLVER:
WEIGHT_DECAY: 0.0001
LR_POLICY: steps_with_decay
BASE_LR: 0.0025
GAMMA: 0.1
MAX_ITER: 60000
STEPS: [0, 30000, 40000]

Equivalent schedules with…

1 GPU:

BASE_LR: 0.0025

MAX_ITER: 60000

STEPS: [0, 30000, 40000]

2 GPUs:

BASE_LR: 0.005

MAX_ITER: 30000

STEPS: [0, 15000, 20000]

4 GPUs:

BASE_LR: 0.01

MAX_ITER: 15000

STEPS: [0, 7500, 10000]

8 GPUs:

BASE_LR: 0.02

MAX_ITER: 7500

STEPS: [0, 3750, 5000]

RPN:
ASPECT_RATIOS: (0.5, 1, 2)
FPN:
FPN_ON: True
MULTILEVEL_ROIS: True
MULTILEVEL_RPN: True
RPN_ASPECT_RATIOS: (0.5, 1, 2)
FAST_RCNN:
ROI_BOX_HEAD: fast_rcnn_heads.add_roi_2mlp_head
ROI_XFORM_METHOD: RoIAlign
ROI_XFORM_RESOLUTION: 7
ROI_XFORM_SAMPLING_RATIO: 2
TRAIN:
WEIGHTS: pretrained_models/ImageNetPretrained/MSRA/R-50.pkl
DATASETS: (‘coco_2014_train’,)
SCALES: (500,)
SNAPSHOT_ITERS: 5000
MAX_SIZE: 833
BATCH_SIZE_PER_IM: 256
RPN_PRE_NMS_TOP_N: 2000 # Per FPN level
TEST:
DATASETS: (‘coco_2014_val’,)
SCALES: (500,)
MAX_SIZE: 833
NMS: 0.5
RPN_PRE_NMS_TOP_N: 1000 # Per FPN level
RPN_POST_NMS_TOP_N: 1000
FORCE_JSON_DATASET_EVAL: True
OUTPUT_DIR: .

主要修改以下内容： 1.MODEL里的NUM_CLASSES，根据自己的数据集进行修改，包括背景。 2.TRAIN和TEST里的DATASETS，设置为如上所示，因为我们已经修改了dataset_catalog.py中的相应内容并指向了我们的数据集，所以我们可以直接使用这两个数据集的名字。 3.TRAIN里的WEIGHTS是一个预训练的模型，如果本地有，修改指向它即可。 4.关于NUM_GPUS，这个需要与SOLVER里面的学习率等一干参数进行联动设置。

3.2 开始训练

设置好数据集和配置文件，即可开始训练模型了。回到$DETECTRON根目录，打开Terminal运行如下命令即可开始训练：

python tools/train_net.py –cfg configs/animals/animals_1gpu_e2e_faster_rcnn_R-50-FPN.yaml OUTPUT_DIR detectron-output

这样即可使用Detectron基于自己的数据集训练各种模型了。

参考资料

Detectron官方设置数据集说明。 Detectron官方GETTING STARTED。使用Detectron训练PASCAL VOC数据集。