资源 | 一个基于PyTorch的目标检测工具箱，商汤联合港中文开源mmdetection

Posted 2021-05-02 机器之心

机器之心整理

参与：路、思

近日，商汤和港中文联合开源了 mmdetection，这是一个基于 PyTorch 的开源目标检测工具包，属于香港中文大学多媒体实验室 open-mmlab 项目的一部分。该工具包支持 Mask RCNN 等多种流行的检测框架，读者可在 PyTorch 环境下测试不同的预训练模型及训练新的检测分割模型。

机器之心Synced 计算机视觉 小程序

根据 GitHub 项目所述，mmdetection 的主要特征可以总结为以下几个方面：

• 模块化设计：你可以通过连接不同组件轻松构建自定义目标检测框架。

• 支持多个框架，开箱即用：该工具包直接支持多种流行的检测框架，如 Faster RCNN、Mask RCNN、RetinaNet 等。

• 高效：所有基础边界框和掩码运算都在 GPU 上运行。不同模型的训练速度大约比 FAIR 的 Detectron 快 5% ~ 20%。

• 当前最优：这是 MMDet 团队的代码库，该团队赢得了 2018 COCO 检测挑战赛的冠军。

除了 mmdetection，研究者们还开源了用于计算机视觉研究的 mmcv 库，mmdetection 很多算法的实现都依赖于 mmcv 库。

此次项目的参与者之一，香港中文大学陈恺博士介绍说：「mmcv 和 mmdetection 同属于 open-mmlab 项目，这个项目会将我们 lab 的一些 research project 有计划、系统性地开源。mmcv 是一个基础库，主要分为两个部分，一部分是和 deep learning framework 无关的一些工具函数，比如 IO/Image/Video 相关的一些操作，另一部分是为 PyTorch 写的一套训练工具，可以大大减少用户需要写的代码量，同时让整个流程的定制变得容易。」

对比 Detection 

据陈恺博士介绍，10 月 12 日商汤正式开源了两个项目，即 mmcv 和 mmdetection：

mmdetection 是基于 MMDet 队伍 COCO 比赛（Detection 2018 winner）时的 codebase 重构，先放一张 test-dev 的结果。

mmdetection 原计划九月底 release，由于某些不可抗力（比如拖延症）延迟到了现在，感觉如释重负。第一个版本中实现了 RPN、Fast R-CNN、Faster R-CNN、Mask R-CNN，近期还计划放出 RetinaNet 和 Cascade R-CNN。（如果由于赶 CVPR 跳票了，请不要介意，介意也不会 release 的。）

先简单介绍一下和 Detectron 的对比

• performance 稍高

• 训练速度稍快

• 所需显存稍小

但更重要的是，基于 PyTorch 和基于 Caffe2 的 code 相比，易用性是有代差的。成功安装 Detectron 的时间，大概可以装好一打的 mmdetection 吧。

当然 Detectron 有些优势也很明显，作为第一个全面的 detection codebase，加上 FAIR 的金字招牌，关注人数和用户很多（虽然吐槽也比较多），release 的模型也比较全面。我们也在努力扩充 model zoo，奈何人力和算力还是有很大差距，所以还需要时间。

具体说说上面提到的三个方面吧。首先是 performance，由于 PyTorch 官方 model zoo 里面的 ResNet 结构和 Detectron 所用的 ResNet 有细微差别（mmdetection 中可以通过 backbone 的 style 参数指定），导致模型收敛速度不一样，所以我们用两种结构都跑了实验，一般来说在 1x 的 lr schedule 下 Detectron 的会高，但 2x 的结果 PyTorch 的结构会比较高。

速度方面 Mask R-CNN 差距比较大，其余的很小。采用相同的 setting，Detectron 每个 iteration 需要 0.89s，而 mmdetection 只需要 0.69s。Fast R-CNN 比较例外，比 Detectron 的速度稍慢。另外在我们的服务器上跑 Detectron 会比官方 report 的速度慢 20% 左右，猜测是 FB 的 Big Basin 服务器性能比我们好？

显存方面优势比较明显，会小 30% 左右。但这个和框架有关，不完全是 codebase 优化的功劳。一个让我们比较意外的结果是现在的 codebase 版本跑 ResNet-50 的 Mask R-CNN，每张卡（12 G）可以放 4 张图，比我们比赛时候小了不少。

除了陈恺博士在知乎上的回答，mmdetection 的 GitHub 项目页面上还有很多实现性能与对比数据：

1. 性能

开发者报告了使用使用 caffe-style 和 pytorch-style ResNet 骨干网络的结果，前者的权重来自 Detectron 中 MODEL ZOO 的预训练模型，后者的权重来自官方 model zoo。

2. 训练速度

训练速度的单位是 s/iter，数值越低代表速度越高。

3. 推断测试

推断速度的单位是 fps (img/s)，数值越高代表效果越好。

测试与训练

首先安装 mmdetection 需要以下环境：

• Linux (tested on Ubuntu 16.04 and CentOS 7.2)

• Python 3.4+

• PyTorch 0.4.1 and torchvision

• Cython

• mmcv

使用预训练模型进行推断

1. 测试数据集

开发者允许在一块 GPU 上运行一或多个进程，如在 8 块 GPU 运行 8 个进程或 16 个进程。当单个进程的 GPU 工作负载并不是很大时，运行多个进程将加速测试，这由参数 --proc_per_gpu <PROCESS_NUM> 指定。

测试和保存结果：

python tools/test.py <CONFIG_FILE> <CHECKPOINT_FILE> --gpus <GPU_NUM> --out <OUT_FILE>

要想执行测试后的评估，你需要添加 --eval <EVAL_TYPES>。支持类型包括：

• proposal_fast：使用 mmdetection 的代码求 proposal 的召回率。（应该可以得到和官方估计同样的结果。）

• proposal：使用 COCO 提供的官方代码求 proposal 的召回率。

• bbox：使用 COCO 提供的官方代码求 box AP 值。

• segm：使用 COCO 提供的官方代码求 mask AP 值。

• keypoints：使用 COCO 提供的官方代码求 keypoint AP 值。

例如，估计使用 8 个 GPU 的 Mask R-CNN，并将结果保存为 results.pkl：

python tools/test.py configs/mask_rcnn_r50_fpn_1x.py <CHECKPOINT_FILE> --gpus 8 --out results.pkl --eval bbox segm

在测试过程中可视化结果同样很方便，只需添加一个参数 --show：

python tools/test.py <CONFIG_FILE> <CHECKPOINT_FILE> --show

2. 测试图像

开发者提供了一些高级 api 来测试图像：

import mmcv
from mmcv.runner import load_checkpoint
from mmdet.models import build_detector
from mmdet.apis import inference_detector, show_result

cfg = mmcv.Config.fromfile('configs/faster_rcnn_r50_fpn_1x.py')
cfg.model.pretrained = None

# construct the model and load checkpoint
model = build_detector(cfg.model, test_cfg=cfg.test_cfg)
_ = load_checkpoint(model, 'https://s3.ap-northeast-2.amazonaws.com/open-mmlab/mmdetection/models/faster_rcnn_r50_fpn_1x_20181010-3d1b3351.pth')

# test a single image
img = mmcv.imread('test.jpg')
result = inference_detector(model, img, cfg)
show_result(img, result)

# test a list of images
imgs = ['test1.jpg', 'test2.jpg']
for i, result in enumerate(inference_detector(model, imgs, cfg, device='cuda:0')):
    print(i, imgs[i])
    show_result(imgs[i], result)

训练模型

mmdetection 使用 MMDistributedDataParallel 和 MMDataParallel 分别实现分布式训练和非分布式训练。

开发者建议在单个机器上也要使用分布式训练，因为它速度更快，而非分布式训练可以用于 debug 或其他目的。

1. 分布式训练

mmdetection 潜在支持多种 launch 方法，如 PyTorch 的内置 launch utility、 slurm 和 MPI。

开发者使用 PyTorch 内置的 launch utility 提供训练脚本：

./tools/dist_train.sh <CONFIG_FILE> <GPU_NUM> [optional arguments]

支持的参数有：

• --validate：训练过程中每 k（默认值为 1）个 epoch 执行估计。

• --work_dir <WORK_DIR>：如果指定，配置文件中的路径将被重写。

2. 非分布式训练

python tools/train.py <CONFIG_FILE> --gpus <GPU_NUM> --work_dir <WORK_DIR> --validate

知乎问答参考链接：https://www.zhihu.com/question/294578141