上一个系列文章 TensorFlow 训练自己的目标检测器 以及 TensorFlow 训练 Mask R-CNN 模型， 说明了怎么用 TensorFlow 开源的目标检测和实例分割接口 TensorFlow/models/research/object_detection 来基于自己的数据训练 Mask R-CNN 模型。但这些都是使用它自带的、而不是自己定义的模型来训练，因此不一定适用自己的数据集，比如，对于一些很简单的数据集，就没必要用它自带的非常复杂的模型。这一篇文章试图弥补这个缺点，来讲述怎么，基于 Tensorflow Detection API，自定义模型训练 Mask R-CNN，相关官方简易文档见 So you want to create a new model!。
作为开始，让我们回顾一下 Mask R-CNN 目标检测与实例分割的过程：对于一张给定的图像，首先经过一个卷积神经网络（称为特征提取器）从图像中提取特征，得到一张特征映射；接着，从这张特征映射预测 2 个分支，分别预测：目标得分、目标边框；以上整个过程便组成了 Mask R-CNN 的第一阶段，整个网络称为候选区域网络（Region Proposal Network， RPN）；其次，从第一阶段预测的结果中筛选出高置信度的目标区域，将这些区域从特征映射中裁剪出来，送入第二阶段的预测网络，更精细的预测 3 个分支：类概率、目标边框、目标实例；最后，通过非极大值抑制等后处理后操作后输出最终的检测结果。从以上过程可见，整个 Mask R-CNN 框架最灵活的地方是特征提取器（不同的卷积神经网络定义不同的特征提取器），而其它地方基本都可以固定不变。因此，借助 Tensorflow Detection API 搭建的 Mask R-CNN 框架，要实现自定义模型训练已变得非常简单，只需要重写一个特征提取器（Feature Extractor）即可。而这可以模仿 models/research/object_detection/models 文件夹的文件来写，这个文件夹内的所有文件都是特征提取器的定义。
本文给出一个简单的示例，编写一个深度较小的特征提取器，嵌入到 Tensorflow Detection API 来实现自定义模型训练 Mask R-CNN。
所有代码见 GitHub: mask_rcnn_customized。
一、自定义特征提取器
要自定义特征提取器，需要重载 models/research/object_detection/meta_architectures 文件夹内的 Mask R-CNN 框架类 faster_rcnn_meta_arch.py 中的特征提取器抽象类 FasterRCNNFeatureExtractor。因此，只需要继承该类，再重新定义该类的抽象函数 process、_extract_proposal_features、_extract_box_classifier_features 即可。我们仿造文件夹 models/research/object_detection/models 内的 faster_rcnn_resnet_v1_feature_extractor.py 来写。
（1） 定义简单卷积模型：ResNet-20
前面已经交代过，特征提取器是一个卷积神经网络，它负责从图像中提取特征，用于后续两阶段的预测。简单起见，我们用 ResNet 的残差模块（Residual Module）来定义一个 20 层的网络，如下（见文件 custom_resnet.py）：
from tensorflow.contrib.slim import nets
resnet_v1_block = nets.resnet_v1.resnet_v1_block
def resnet_v1_20(inputs,
num_classes=None,
is_training=True,
global_pool=True,
output_stride=None,
spatial_squeeze=True,
store_non_strided_activations=False,
reuse=None,
scope='resnet_v1_20'):
"""ResNet-20 model. See resnet_v1() for arg and return description."""
blocks = [
resnet_v1_block('block1', base_depth=64, num_units=1, stride=2),
resnet_v1_block('block2', base_depth=128, num_units=1, stride=2),
resnet_v1_block('block3', base_depth=256, num_units=1, stride=2),
resnet_v1_block('block4', base_depth=512, num_units=3, stride=1)
]
return nets.resnet_v1.resnet_v1(
inputs,
blocks,
num_classes,
is_training,
global_pool=global_pool,
output_stride=output_stride,
include_root_block=True,
reuse=reuse,
scope=scope)
一个残差模块由 3 个卷积层组成（见下图-右），在 TensorFlow 的实现里称为一个 Unit，多个 Unit 的组合称为一个 Block。如上面的代码，使用了 4 个 Block，它们的 Unit 个数（num_units）分别为 1,1,1,3，因此总共有 20 = （1 + 1 + 1 + 3） x 3 + 1 + 1 个卷积层，最后的 1 + 1 指的分别是卷积核为 11 x 11 的网络第一个卷积层和最后的 softmax 输出层。
（2）重载特征提取器
接下来，来写自定义的 ResNet-20 对应的特征提取器，代码如下（见 custom_faster_rcnn_resnet_v1_feature_extractor.py）：
# -*- coding: utf-8 -*-
"""
Created on Thu Nov 1 14:18:07 2018
@author: shirhe-lyh
ResNet V1 Faster R-CNN customized implementation.
"""
import tensorflow as tf
from tensorflow.contrib.slim import nets
from object_detection.meta_architectures import faster_rcnn_meta_arch
from object_detection.models import custom_resnet
slim = tf.contrib.slim
resnet_v1_block = nets.resnet_v1.resnet_v1_block
class CustomFasterRCNNResnetV1FeatureExtractor(
faster_rcnn_meta_arch.FasterRCNNFeatureExtractor):
"""Faster R-CNN ResNet v1 feature extractor customized implementation."""
def __init__(self,
architecture,
resnet_model,
is_training,
first_stage_features_stride,
batch_norm_trainable=False,
reuse_weights=None,
weight_decay=0.0):
"""Constructor.
Args:
architecture: Architecture name of the ResNet V1 model.
resnet_model: Definition of the ResNet V1 model.
is_training: See base class.
batch_norm_trainable: See base class.
first_stage_features_stride: See base class.
batch_norm_trainable: See base class.
reuse_weights: See base class.
weight_decay: See base class.
Raises:
ValueError: If `first_stage_features_stride` is not 8 or 16.
"""
if first_stage_features_stride != 8 and first_stage_features_stride !=16:
raise ValueError('`first_stage_features_stride` must be 8 or 16.')
self._architecture = architecture
self._resnet_model = resnet_model
super(CustomFasterRCNNResnetV1FeatureExtractor, self).__init__(
is_training, first_stage_features_stride, batch_norm_trainable,
reuse_weights, weight_decay)
def preprocess(self, resized_inputs):
"""Faster R-CNN ResNet V1 preprocessing.
Args:
resized_inputs: A [batch, height_in, width_in, channels] float32
tensor representing a batch of images with values between 0
and 255.0.
Returns:
preprocessed_inputs: A [batch, height_out, width_out, channels]
float32 tensor representing a batch of images.
"""
channel_means = [123.68, 116.779, 103.939]
return resized_inputs - [[channel_means]]
def _extract_proposal_features(self, preprocessed_inputs, scope):
"""Extracts first stage RPN features.
Args:
preprocessed_inputs: A [batch, height, width, channels] float32
tensor representing a batch of images.
scope: A scope name.
Returns:
rpn_feature_map: A tensor with shape [batch, height, width, depth].
activations: A dictionary mapping feature extractor tensor names
to tensors.
Raises:
InvalidArgumentError: If the spatial size of `preprocessed_inputs`
(height or width) is less than 33.
ValueError: If the created network is missing the required
activation.
"""
if len(preprocessed_inputs.get_shape().as_list()) != 4:
raise ValueError('`preprocessed_inputs` must be 4 dimensional, '
'got a tensor of shape %s' %
preprocessed_inputs.get_shape())
shape_assert = tf.Assert(
tf.logical_and(
tf.greater_equal(tf.shape(preprocessed_inputs)[1], 33),
tf.greater_equal(tf.shape(preprocessed_inputs)[2], 33)),
['image size must at least be 33 in both height and width.'])
with tf.control_dependencies([shape_assert]):
# Disables batchnorm for fine-tuning with smaller batch sizes.
# TODO(chensun): Figure out if it is needed when image
# batch size is bigger.
with slim.arg_scope(nets.resnet_utils.resnet_arg_scope(
batch_norm_epsilon=1e-5,
batch_norm_scale=True,
weight_decay=self._weight_decay)):
with tf.variable_scope(self._architecture,
reuse=self._reuse_weights) as var_scope:
_, activations = self._resnet_model(
preprocessed_inputs,
num_classes=None,
is_training=self._train_batch_norm,
global_pool=False,
output_stride=self._first_stage_features_stride,
spatial_squeeze=False,
scope=var_scope)
handle = scope + '/%s/block3' % self._architecture
return activations[handle], activations
def _extract_box_classifier_features(self, proposal_feature_maps, scope):
"""Extracts second stage box classifier features.
Args:
proposal_feature_maps: A 4-D float tensor with shape [batch_size *
self.max_num_proposals, crop_height, crop_width, depth]
representing the feature map croped to each proposal.
scope: A scope name (unused).
Returns:
proposal_classifier_features: A 4-D float tensor with shape
[batch_size * self.max_num_proposals, height, width, depth]
representing box classifier features for each proposal.
"""
with tf.variable_scope(self._architecture, reuse=self._reuse_weights):
with slim.arg_scope(nets.resnet_utils.resnet_arg_scope(
batch_norm_epsilon=1e-5,
batch_norm_scale=True,
weight_decay=self._weight_decay)):
with slim.arg_scope([slim.batch_norm],
is_training=self._train_batch_norm):
blocks = [
nets.resnet_utils.Block(
'block4', nets.resnet_v1.bottleneck,
[{'depth': 2048,
'depth_bottleneck': 512,
'stride': 1
}] * 3)
]
proposal_classifier_features = (
nets.resnet_utils.stack_blocks_dense(
proposal_feature_maps, blocks))
return proposal_classifier_features
class CustomFasterRCNNResnet20FeatureExtractor(
CustomFasterRCNNResnetV1FeatureExtractor):
"""Faster R-CNN ResNet V1 20 feature extractor implementation."""
def __init__(self,
is_training,
first_stage_features_stride,
batch_norm_trainable=False,
reuse_weights=None,
weight_decay=0.0):
"""Construtor.
Args:
is_training: See base class.
first_stage_features_stride: See base class.
batch_norm_trainable: See base class.
reuse_weights: See base class.
weight_decay: See base class.
Raises:
ValueError: If `first_stage_features_stride` is not 8 or 16, or
if `architecture` is not supported.
"""
super(CustomFasterRCNNResnet20FeatureExtractor, self).__init__(
'resnet_v1_20', custom_resnet.resnet_v1_20, is_training,
first_stage_features_stride, batch_norm_trainable,
reuse_weights, weight_decay)
因为我们自定义的模型恰好也是 ResNet，因此基本照抄了官方示例文件 faster_rcnn_resnet_v1_feature_extractor.py 的内容，除了类名和所用卷积模型不同之外。
重载特征提取器，需要按照官方的规定，一方面要重载 3 个函数 process、_extract_proposal_features、_extract_box_classifier_features，另一方面，还要遵循约定：_extract_proposal_features 返回的应该是自定义模型的 中间 某个卷积层的结果，而剩下的卷积层由 _extract_box_classifier_features 返回。比如，对于我们自定义的 ResNet-20，_extract_proposal_features 函数只返回前 3 个 Block，最后的第 4 个 Block 由 _extract_box_classifier_features 返回。官方的 Remark 如下：
最后的类 CustomFasterRCNNResnet20FeatureExtractor 就是我们自定义的特征提取器，要使用这个类来训练 Mask R-CNN，还需要将它加入到 TensorFlow Object Detection API 框架里，即需要对它进行注册。
还需要注意的一点是：_extract_proposal_features 这个函数要求输入的图像分辨率必须不小于 33 x 33，因对图像做预处理（比如对图像做缩放）后不要违背了这个硬性要求。
二、模型注册
模型注册非常简单：首先将自定义的文件 custom_resnet.py 和 custom_faster_rcnn_resnet_v1_feature_extractor.py 复制到文件夹 models/research/object_detection/models 里，然后修改文件夹 models/research/object_detection/builders 里的文件 model_builder.py，在模块导入最后加入一条导入语句：
from object_detection.models import \
custom_faster_rcnn_resnet_v1_feature_extractor as custom_frcnn_resnet_v1
之后在该文件的 FASTER_RCNN_FEATURE_EXTRACTOR_CLASS_MAP 字典里加入 key-value 对：
'custom_frcnn_resnet20':
custom_frcnn_resnet_v1.CustomFasterRCNNResnet20FeatureExtractor,
自此，模型注册就完成了。model_builder.py 改动后的文件也也已经上传到 GitHub: mask_rcnn_customized。接下来，只需要准备数据启动训练就可以验证我们前面的工作是否成功了。
三、数据准备
GitHub: mask_rcnn_customized 这个项目的里的文件 shape_mask_generator.py 文件用来生成简单的正多边形几何形状，我们以此来生成训练数据。简单起见，我们以生成等边三角形和正方形，且每幅图像有且仅有一个对象为例：
运行项目里的文件 generate_datasets.py（GitHub 项目内已经生成了数据，此步可跳过；如果要重新生成，请删除 datasets 文件夹）:
$ python3 generate_datasets.py
会在当前路径下生成一个 datasets 的文件夹，里面包含 5000 张图像（images 文件夹）和对应的掩模（masks 文件夹），另外还有一个记录图像与掩模对应关系，以及记录每个目标的 boundingbox 的标注文件：annotations.json。如果，你想生成更多的图像，请修改文件 generate_datasets.py 的参数 num_samples。因为，每张掩模都是 0-1 二值的灰度图，所以我们直接观看 masks 文件夹里面的 .png 图像时都是全黑的。如果要将其 mask 显示出来，请执行：
$ python3 visualize_masks.py
然后到 /datasets/masks_recgonized 文件夹内查看。
图像和掩模生成好后，需要将它们写入 TFRecord 文件，执行（GitHub 项目内已经生成了 .record 文件，此步可跳过；如果要重新生成，请直接执行）：
$ python3 generate_tfrecord.py
在 datasets 文件夹内生成 train.record 和 val.record 文件，下面，就可以开始训练了。
四、模型训练
mask_rcnn_customized 里已经配置好了类名与类标号转化文件 shape_label_map.pbtxt 以及模型参数配置文件 mask_rcnn_customized_resnet_v1_shape.config（需要修改路径：train_input_reader: {...} 中的 input_path 和 label_map_path，以及 eval_input_reader: {...} 中的 input_path 和 label_map_path），要启动训练，进入你配置好的 TensorFlow Object Detection API 项目的文件夹 models/research/object_detection 内，执行：
$ python3 model_main.py \
--model_dir Path/to/mask_rcnn_customized/training \
--pipeline_config_path Path/to/mask_rcnn_customized_resnet_v1_shape.config
即可。因为要使用自己定义的特征提取器，在配置 .config 文件时必须要将其中的
model {
faster_rcnn {
feature_extractor {
type: 'custom_frcnn_resnet20'
...
}
...
}
...
}
type 字段修改为自己在 二 中注册特征提取器时的特征提取器名。
训练开始后，使用：
$ tensorboard --logdir Path/to/mask_rcnn_customized/training
命令得到浏览器链接，打开该链接可实时监督训练过程。训练过程的精度/损失曲线大致如下：
训练结束后，在路径 models/research/object_detection 下执行：
$ python3 export_inference_graph.py \
--trained_checkpoint_prefix Path/to/training/model.ckpt-40000 \
--output_directory Path/to/converted_pb_file_saving_directory \
--pipeline_config_path Path/to/mask_rcnn_customized_resnet_v1_shape.config
将训练保存的 .ckpt 模型转化为 .pb 格式，方便后续调用。其中，参数 trained_checkpoint_prefix 指定训练后 .ckpt 模型保存的路径（详细指定到某个训练次数时的模型），output_directory 指定转化后的 .pb 格式模型的保存路径（填写某个路径下的文件夹，该文件夹可以不存在，比如填写 /home/.../mask_rcnn_customized/training/frozen_inference_graph_pb），pipeline_config_path 指定 mask_rcnn_customized_resnet_v1_shape.config 文件路径。
五、结果展示
当顺利完成以上所有步骤之后，就可以运行：
$ python3 predict.py
来进行预测了（注意：需要将 predict.py 文件中的 PATH_TO_CKPT 填写为上一步转化来的 .pb 文件所在的路径。如果，执行 python3 export_inference_graph.py ... 命令时，你填写的 output_directory 是 /home/.../mask_rcnn_customized/training/frozen_inference_graph_pb，则使用默认路径而不需要修改）。执行后，会在当前路径下生成 test_images 文件夹，里面会输出 10 张测试图像，以及他们的检测结果：
由于只使用了 5000 个训练样本，且没有仔细调参，可以发现模型对于三角形的检测和分割效果不是太理想。如果，你想改善效果，可以使用更多的训练数据。
说明：
1.如果不使用预训练模型，比如我们自定义的 ResNet-20 就找不到训练好的参数，在修改配置文件 xxx.config 时要将 train_config: {...} 里的其中两行：
#fine_tune_checkpoint: "PATH_TO_BE_CONFIGURED/model.ckpt"
#from_detection_checkpoint: true
注释掉，这样所有模型参数都会随机初始化。
2.训练过程如果报如下错误：TypeError: can't pickle dict_values objects，则将 models/research/object_detection/model_lib.py 中第 418 行的 category_index.values() 改成 list(category_index.values()) 即可