专栏首页机器之心资源 | 一个基于PyTorch的目标检测工具箱,商汤联合港中文开源mmdetection

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

机器之心整理

参与:路、思

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

项目地址:https://github.com/open-mmlab/mmdetection

根据 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

本文分享自微信公众号 - 机器之心(almosthuman2014)

原文出处及转载信息见文内详细说明,如有侵权,请联系 yunjia_community@tencent.com 删除。

原始发表时间:2018-10-17

本文参与腾讯云自媒体分享计划,欢迎正在阅读的你也加入,一起分享。

我来说两句

0 条评论
登录 后参与评论

相关文章

  • 你要「老婆」不要?谷歌程序员20行代码把二次元「老婆」带到现实世界

    还记得前不久的「‍东升发型生成器」吗?在那个视频中,B 站 up 主 Mars 隆重推出了自制的秃头生成器。不管你现在有多年轻,头发有多么浓密,那个生成器都能让...

    机器之心
  • 登上Nature,DeepMind新论文提前48小时预测急性肾损伤

    DeepMind 近日在 Nature 发表论文,发布了一种新的人工智能算法,可以提前两天预测到一种可避免的病变。通过和美国退伍军人事务部一起合作,研究人员开发...

    机器之心
  • 学界 | 优于VAE,为万能近似器高斯混合模型加入Wasserstein距离

    使用生成式隐变量模型的无监督学习提供了一种强大且通用的方法来从大型无标签数据集中学习潜在的低维结构。通常训练该模型的两种最常见的技术是变分自编码器(VAE)[1...

    机器之心
  • 构建另一种服务网格:使用SMI规范的新方法

    在这篇文章中,我将通过深入研究Maesh项目背后的技术细节,探究服务网格接口(Service Mesh Interface,SMI)规范的高级概念,是什么使该项...

    CNCF
  • 献给即将35岁的初学者,焦虑 or 出路?

    一直以来,都有许多人说“程序员或测试员是个吃青春饭的职业”,甚至还有说“35 岁混不到管理就等于失业”的言论。

    测试开发技术
  • AI变身记:不光能有人的智能,还要像狗一样“思考”

    用户1737318
  • 无需标注数据集,自监督注意力机制就能搞定目标跟踪

    想要了解什么是自监督注意力机制,我们可能需要先去了解什么是光流估计(optical flow estimation),以及它为何被人类和计算机视觉系统作为一种目...

    机器之心
  • 原创 | 算法工程师为什么成天做数据,都做哪些数据?

    大家好,前几天群里有小伙伴说希望看到更多的算法工程师的日常。其实对于算法工程师而言,最大的日常就是做数据了,所以给大家分享一下做数据的那些事。

    TechFlow-承志
  • 如何构建一个好的数据挖掘模型

    1 首先要保证数据质量,特征值需要选择好,其次需要弄明白业务需求,确立你的分析目标是分类还是预测

    许卉
  • NuGet 更新库新版本的缓存问题

    我有一个 NuGet 库有新的版本,但是我的服务器速度不够快,此时我第一次使用 NuGet 还原找不到库。在我服务器索引完成之后,再次使用 NuGet 会依然找...

    林德熙

扫码关注云+社区

领取腾讯云代金券