物体检测和分割轻松上手：从detectron2开始（下篇）

AI编辑：我是小将

Model使用和创建

模型是目标检测和分割最核心的部分，detectron2的模型也是模块的，主要包括4个核心的部分：

• backbone：模型的CNN特征提取器，目前只支持resnet，另外一点是detectron2也把FPN作为backbone的一部分；

• proposal_generator：候选框生成器，目前只支持RPN，一般用于faster R-CNN的一部分，其实RPN单拿出来也是一个简单的one-stage检测模型；

• roi_heads：faster R-CNN系列的detector，包括ROIPooler，box_head，mask_head等，其中ROIPooler就是指的ROIPool和ROIAlign方法；

• meta_arch：定义最终的模型，不能说是一个单独的模块，应该要集成backbone，proposal_generator和roi_heads构建最终模型。

模块的话好处是可以复用模块来进行组合，比如你可以组合backbone，proposal_generator和roi_heads来构建不同的模型。由于detectron2官方只是支持RCNN模型，所以实现的模块可能不够通用。YOLOv4中给出了更加通用的模块划分：

这个划分更加通用化，主流的检测模型都囊括其中。目前mmdetection框架是按照类似的模块划分来设计的，所以支持更多的模型。

前面已经提过，模型构建是用统一的接口：

from detectron2.modeling import build_model
model = build_model(cfg)  # 得到的是一个torch.nn.Module

不论是使用模型还是要自定义模型，必须要了解detectron2中的模型的输入和输出格式。模型的输入是一个list[dict]，每个dict是一个样本的图像以及标注信息，具体如下：

• “image”: (C, H, W)格式的Tensor，cfg.INPUT.FORMAT决定图像的channel格式，默认是BGR格式，模型侧会采用cfg.MODEL.PIXEL_{MEAN,STD}对image进行归一化；

• “instances”: Instances实例，包括这些字段：gt_boxes，gt_classes，gt_masks，gt_keypoints，不同的模型需求不一样；

• “proposals”：Instances实例，RPN产生的结果，用于Faster R-CNN等模型输入，包括proposal_boxes和objectness_logits字段；

• “height”, “width”：模型最终输出时的图像大小，不一定和image的size一致，比如你做了resize，但是模型最终可以还原到原来的大小；

• “sem_seg”: (H, W)格式的Tensor，用于语义分割GT，类别从0开始（其实detectron是支持分割的，这是一个亮点）；

模型的输入主要包括上述所示的字段，但是根据模型的不同，所需要的字段也有变化。前面说过mapper的输出是作为模型的输入，所以mapper的输出格式也要满足上述要求，在dataloader部分已经给出了实例分割的case。

模型在训练时输出的是dict[str->ScalarTensor]，即各个部分的loss用于模型训练，重点是测试的输出格式，模型输出也是一个list[dict]，每个image对应一个dict：

• “instances”: Instances实例，包括:pred_boxes， scores， pred_classes，pred_masks，pred_keypoints，也是模型不同包含字段有不同；

• “sem_seg”: (num_categories, H, W)格式的Tensor，语义分割预测结果；

• “proposals”：Instances实例，包括:proposal_boxes和objectness_logits，RPN的输出；

• “panoptic_seg”: (Tensor, list[dict])，全景分割结果；

同样地，不同类型的模型其输出字段也是不同的，如实例分割的输出格式如下：

{'instances': Instances(num_instances=15, image_height=480, image_width=640, fields=[pred_boxes: Boxes(tensor), scores: tensor, pred_classes: tensor, pred_masks: tensor])}

更进一步地，我们有时候需要自定义新的模型，这就要利用前面所说的detectron2注册机制。由于detectron2是模块化设计的，你也可以自定义模型的某一个模块，比如你想实现一个新的backbone，那么简单的case如下：

from detectron2.modeling import BACKBONE_REGISTRY, Backbone, ShapeSpec

## 注册backbone
@BACKBONE_REGISTRY.register()
class ToyBackBone(Backbone):
  def __init__(self, cfg, input_shape):
    super().__init__()
    # create your own backbone
    self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=16, padding=3)

  # 输出格式是dict，这是detectron2的约定格式，与后面的detector相适应
  def forward(self, image):
    return {"conv1": self.conv1(image)}

  # 这个方法是必须的，因为detector的head需要backbone的输出channel
  def output_shape(self):
    return {"conv1": ShapeSpec(channels=64, stride=16)}

使用这个新的backbone，只需要在配置文件中修改cfg.MODEL.BACKBONE.NAME = 'ToyBackBone'，那么build_model(cfg)得到的模型将使用这个新的backbone。一个要注意的点是，自定义创建的backbone要符合一定规范，这方面可以参考mobilenetv2和vovnet这两个backbone的自定义创建。

其实要学习如何创建自己的模型，一方面首先要掌握的detectron2中核心模型的实现，另外一方面detectron2下的projects也提供了很多优质的模型供参考学习。

模型训练和评测

我们只需要用engine里的DefaultTrainer就可以进行完成训练。前面已经构建了balloon数据集，这里采用mask_rcnn_R_50_FPN模型进行训练，为了加快收敛，使用COCO数据集上训练的模型进行finetune：

from detectron2.·engine import DefaultTrainer
from detectron2.config import get_cfg

cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
cfg.DATASETS.TRAIN = ("balloon_train",)
cfg.DATASETS.TEST = ()
cfg.DATALOADER.NUM_WORKERS = 2
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")  # COCO模型上finetune
cfg.SOLVER.IMS_PER_BATCH = 2
cfg.SOLVER.BASE_LR = 0.00025  
cfg.SOLVER.MAX_ITER = 300    # 迭代步数
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128 
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1  # 只有一类(ballon)

os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
trainer = DefaultTrainer(cfg) 
trainer.resume_or_load(resume=False)
trainer.train()

运行之后，将可以看到模型信息，数据集信息，以及训练过程中的log。训练完成后，可以使用DefaultPredictor对图片进行测试：

cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "model_final.pth")
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.7   # set the testing threshold for this model
cfg.DATASETS.TEST = ("balloon_val", )
predictor = DefaultPredictor(cfg)

from detectron2.utils.visualizer import ColorMode
dataset_dicts = get_balloon_dicts("balloon/val")
for d in random.sample(dataset_dicts, 3):    
    im = cv2.imread(d["file_name"])
    outputs = predictor(im)
    v = Visualizer(im[:, :, ::-1],
                   metadata=balloon_metadata, 
                   scale=0.8, 
                   instance_mode=ColorMode.IMAGE_BW   # remove the colors of unsegmented pixels
    )
    v = v.draw_instance_predictions(outputs["instances"].to("cpu"))
    cv2_imshow(v.get_image()[:, :, ::-1])

如果要对训练好的模型进行评测，就是计算mAP值，只需要：

from detectron2.evaluation import COCOEvaluator, inference_on_dataset
from detectron2.data import build_detection_test_loader
evaluator = COCOEvaluator("balloon_val", cfg, False, output_dir="./output/")
val_loader = build_detection_test_loader(cfg, "balloon_val")
inference_on_dataset(trainer.model, val_loader, evaluator)

------------------------------------------
[06/07 09:30:02 d2.evaluation.coco_evaluation]: Evaluation results for segm: 
|   AP   |  AP50  |  AP75  |  APs  |  APm   |  APl   |
|:------:|:------:|:------:|:-----:|:------:|:------:|
| 77.251 | 85.083 | 84.743 | 5.959 | 60.356 | 93.479 |

DefaultTrainer里面已经包含了学习速率策略，日志，模型保存以及评估等配置，所以对于大多数情况，直接使用DefaultTrainer训练自己的模型即可。detectron2中提供了tools/train_net.py对DefaultTrainer进行了简单的包装，我们可以执行下列命令进行训练：

./train_net.py --config-file ../configs/COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_1x.yaml --num-gpus $N  

对于配置参数，除了直接修改yaml文件后，你也可以在执行命令时修改配置参数，比如调整学习速率：

./train_net.py --config-file ../configs/COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_1x.yaml --num-gpus $N SOLVER.BASE_LR 0.0025

如果你想自定义某些训练逻辑，就可以按照train_net.py那样重写DefaultTrainer的某些方法，比如你希望使用自定义的mapper进行训练，那么只需要重写build_train_loader方法：

@classmethod
def build_train_loader(cls, cfg):
    return build_detection_train_loader(cfg, mapper=your_mapper)

如果你觉得DefaultTrainer满足不了自己的训练需求，更进一步地你可以参考tools/plain_train_net.py实现自己的某些策略。

模型部署

模型部署也是很重要的一环，目前detectron2只支持转化成caffe2，官方提供了转换脚本：tools/deploy/caffe2_converter.py。当然转成ONNX也是可以的，但是并不是所有的op都支持，这个坑就要自己踩了。另外一个很赞的工作是，detectron2作者提供了一种支持转换成TF模型的方法，那就是通过tensorpack FasterRCNN，目前支持Faster R-CNN和Mask R-CNN等模型。

总之，detectron2是一个非常好的目标检测和分割的开源项目，无论是代码质量还是框架灵活性都是上乘。唯一的缺点是就是相比mmdetection，所支持的模型较少，好在可以扩展。

参考

1. facebookresearch/detectron2 https://github.com/facebookresearch/detectron2

2. Detectron2 Beginner's Tutorial https://colab.research.google.com/drive/16jcaJoc6bCFAQ96jDe2HwtXj7BMD_-m5

---

如果觉得不错，请关注一下公众号，也请为小编点个在看，回复“加群”即可进群学习！

推荐阅读

VoVNet：实时目标检测的新backbone网络

想读懂YOLOV4，你需要先了解下列技术(一)

想读懂YOLOV4，你需要先了解下列技术(二)

PyTorch分布式训练简明教程

机器学习算法工程师

---

                                            一个用心的公众号

---