业界 | Facebook开源Mask R-CNN的PyTorch 1.0基准,比mmdetection更快、更省内存

选自GitHub

机器之心编辑

参与:刘晓坤

近日,Facebook AI Research 开源了 Faster R-CNN 和 Mask R-CNN 的 PyTorch 1.0 实现基准:MaskRCNN-Benchmark。相比 Detectron 和 mmdetection,MaskRCNN-Benchmark 的性能相当,并拥有更快的训练速度和更低的 GPU 内存占用。

项目地址:https://github.com/facebookresearch/maskrcnn-benchmark

MaskRCNN-Benchmark 目标检测示例。

Detectron 和 mmdetection

Detectron 是 Facebook AI Research 实现顶尖目标检测算法(包括 Mask R-CNN)的软件系统。该系统是基于 Python 和深度学习框架 Caffe 2 而构建的。Detectron 目前包含以下目标检测算法的实现:Mask R-CNN 、RetinaNet、Faster R-CNN、RPN、Fast R-CNN、R-FCN。

mmdetection 是商汤和港中文近日联合开源的基于 PyTorch 的开源目标检测工具包。该工具包支持 Mask R-CNN 等多种流行的检测框架,读者可在 PyTorch 环境下测试不同的预训练模型及训练新的检测分割模型。和 Detectron 对比,mmdetection 的性能稍高、训练速度稍快、所需显存稍小。

mmdetection 第一个版本中实现了 RPN、Fast R-CNN、Faster R-CNN、Mask R-CNN,近期还计划放出 RetinaNet 和 Cascade R-CNN。但更重要的是,基于 PyTorch 和基于 Caffe2 的 code 相比,易用性是有代差的。成功安装 Detectron 的时间,大概可以装好一打的 mmdetection。

MaskRCNN-Benchmark 项目亮点:

  • PyTorch 1.0:相当或者超越 Detectron 准确率的 RPN、Faster R-CNN、Mask R-CNN 实现;
  • 非常快:训练速度是 Detectron 的两倍,是 mmdection 的 1.3 倍。
  • 节省内存:在训练过程中使用的 GPU 内存比 mmdetection 少大约 500MB;
  • 使用多 GPU 训练和推理;
  • 批量化推理:可以在每 GPU 每批量上使用多张图像进行推理;
  • 支持 CPU 推理:可以在推理时间内于 CPU 上运行。
  • 提供几乎所有参考 Mask R-CNN 和 Faster R-CNN 配置的预训练模型,具有 1x 的 schedule。

MaskRCNN-Benchmark Model Zoo 基线模型性能数据

地址:https://github.com/facebookresearch/maskrcnn-benchmark/blob/master/MODEL_ZOO.md

硬件

8 NVIDIA V100 GPUs

软件:

  • PyTorch version: 1.0.0a0+dd2c487
  • CUDA 9.2
  • CUDNN 7.1
  • NCCL 2.2.13-1

端到端 Mask R-CNN 和 Faster R-CNN 基线模型

所有的基线模型都使用了和 Detectron 相同的实验设置,检测模型权重使用 Caffe2 中的 ImageNet 权重初始化,这和 Detectron 是一样的。预训练模型通过下表中的 model id 链接获取。

和 Detectron、mmdetection 的性能对比

训练速度

下表中的数据单位是秒/迭代,越低越好。(mmdetection 中备注的硬件和 maskrcnn_benchmark 是不同的)

训练内存(越低越好)

推理准确率(越高越好)

Webcam 和 Jupyter notebook demo

该项目提供了一个简单的 webcam demo,展示如何使用 maskrcnn_benchmark 进行推理:

cd demo# by default, it runs on the GPU# for best results, use min-image-size 800
python webcam.py --min-image-size 800# can also run it on the CPU
python webcam.py --min-image-size 300 MODEL.DEVICE cpu# or change the model that you want to use
python webcam.py --config-file ../configs/caffe2/e2e_mask_rcnn_R_101_FPN_1x_caffe2.py --min-image-size 300 MODEL.DEVICE cpu# in order to see the probability heatmaps, pass --show-mask-heatmaps
python webcam.py --min-image-size 300 --show-mask-heatmaps MODEL.DEVICE cpu

安装

教程地址:https://github.com/facebookresearch/maskrcnn-benchmark/blob/master/INSTALL.md

安装要求:

  • PyTorch 1.0 的每日测试版本,安装说明:https://pytorch.org/get-started/locally/
  • torchvision
  • cocoapi
  • yacs
  • (可选)OpenCV(用于 webcam demo)
# maskrnn_benchmark and coco api dependencies
pip install ninja yacs cython
# follow PyTorch installation in https://pytorch.org/get-started/locally/# we give the instructions for CUDA 9.0
conda install pytorch-nightly -c pytorch
# install torchvisioncd ~/github
git clone git@github.com:pytorch/vision.gitcd vision
python setup.py install
# install pycocotoolscd ~/github
git clone git@github.com:cocodataset/cocoapi.gitcd cocoapi/PythonAPI
python setup.py build_ext install
# install PyTorch Detectioncd ~/github
git clone git@github.com:facebookresearch/maskrcnn-benchmark.gitcd maskrcnn-benchmark# the following will install the lib with# symbolic links, so that you can modify# the files if you want and won't need to# re-build it
python setup.py build develop

通过几行代码进行推理

该项目提供了一个 helper 类来简化编写使用预训练模型进行推理的流程,只要在 demo 文件夹下运行以下代码:

from maskrcnn_benchmark.config import cfgfrom predictor import COCODemo

config_file = "../configs/caffe2/e2e_mask_rcnn_R_50_FPN_1x_caffe2.yaml"# update the config options with the config file
cfg.merge_from_file(config_file)# manual override some options
cfg.merge_from_list(["MODEL.DEVICE", "cpu"])

coco_demo = COCODemo(
    cfg,
    min_image_size=800,
    confidence_threshold=0.7,
)# load image and then run prediction
image = ...
predictions = coco_demo.run_on_opencv_image(image)

在 COCO 数据集上执行训练

为了运行以下示例,你首先需要安装 maskrcnn_benchmark。你还需要下载 COCO 数据集,推荐按以下方式符号链接 COCO 数据集的路径到 datasets/。我们使用来自 Detectron 的 GitHub 的 minival 和 valminusminival 集合。

# symlink the coco datasetcd ~/github/maskrcnn-benchmark
mkdir -p datasets/coco
ln -s /path_to_coco_dataset/annotations datasets/coco/annotations
ln -s /path_to_coco_dataset/train2014 datasets/coco/train2014
ln -s /path_to_coco_dataset/test2014 datasets/coco/test2014
ln -s /path_to_coco_dataset/val2014 datasets/coco/val2014

你也可以配置你自己的到数据集的路径。为此,你需要让 maskrcnn_benchmark/config/paths_catalog.py 指向你的数据集保存的位置。你也可以创建一个新的 paths_catalog.py 文件,其实现了相同的两个类,并在训练过程中将它作为一个配置参数 PATHS_CATALOG 传递。

单 GPU 训练

python /path_to_maskrnn_benchmark/tools/train_net.py --config-file "/path/to/config/file.yaml"

多 GPU 训练

该项目使用内部的 torch.distributed.launch 以启动多 GPU 训练。这个来自 PyTorch 的效用函数可以产生我们想要使用 GPU 数目的 Python 进程,并且每个 Python 进程只需要使用一个 GPU。

export NGPUS=8
python -m torch.distributed.launch --nproc_per_node=$NGPUS /path_to_maskrcnn_benchmark/tools/train_net.py --

添加你自己的数据集

该项目添加了对 COCO 类型数据集的支持,为在新数据集上训练添加支持可以通过以下方法实现:

from maskrcnn_benchmark.structures.bounding_box import BoxList
class MyDataset(object):
    def __init__(self, ...):
        # as you would do normallydef __getitem__(self, idx):
        # load the image as a PIL Image
        image = ...# load the bounding boxes as a list of list of boxes# in this case, for illustrative purposes, we use# x1, y1, x2, y2 order.
        boxes = [[0, 0, 10, 10], [10, 20, 50, 50]]
        # and labels
        labels = torch.tensor([10, 20])

        # create a BoxList from the boxes
        boxlist = Boxlist(boxes, size=image.size, mode="xyxy")
        # add the labels to the boxlist
        boxlist.add_field("labels", labels)

        if self.transforms:
            image, boxlist = self.transforms(image, boxlist)

        # return the image, the boxlist and the idx in your datasetreturn image, boxlist, idx

    def get_img_info(self, idx):
        # get img_height and img_width. This is used if# we want to split the batches according to the asp

就这样。你可以添加额外的字段到 boxlist,例如 segmentation masks(使用 structures.segmentation_mask.SegmentationMask),或甚至是你自己的实例类型。如果想了解 COCO 数据集实现的完整过程,可以查看:https://github.com/facebookresearch/maskrcnn-benchmark/blob/master/maskrcnn_benchmark/data/datasets/coco.py

原文发布于微信公众号 - 机器之心(almosthuman2014)

原文发表时间:2018-10-25

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

发表于

我来说两句

0 条评论
登录 后参与评论

相关文章

来自专栏机器之心

教程 | 如何快速训练免费的文本生成神经网络

30750
来自专栏数据派THU

教你用Keras和CNN建立模型识别神奇宝贝!(附代码)

在今天博客的最后,你将会了解如何在你自己的数据库中建立、训练并评估一个卷积神经网络。

79310
来自专栏AI研习社

你在数据预处理上花费的时间,是否比机器学习还要多?

Nuts-ml 是一个新的 Python 数据预处理库,专门针对视觉领域的 GPU 深度学习应用。 它以独立、可复用的单元模块的形式,提供主流数据预处理函数。...

41280
来自专栏专知

【前沿】TensorFlow Pytorch Keras代码实现深度学习大神Hinton NIPS2017 Capsule论文

【导读】10月26日,深度学习元老Hinton的NIPS2017 Capsule论文《Dynamic Routing Between Capsules》终于在a...

44160
来自专栏专知

【实战】最新Deep Learning with Keras图书加代码,教你从零开发一个复杂深度学习模型(附下载)

Deep Learning with Keras 这个代码库是《Deep Learning with Keras》的配套代码。它包含全书从头到尾所有所需要用到的...

60650
来自专栏ATYUN订阅号

不用P图软件,TensorFlow pix2pix也能帮你打造出六块腹肌!

可视化结果是十分强大的。然而,在健身领域,要想清楚地看到未来的锻炼结果往往是很困难的。我们是否可以利用深度学习让人们更接近他们的个人健康目标,从而帮助他们设想未...

65040
来自专栏机器之心

教程 | face-api.js:在浏览器中进行人脸识别的JavaScript接口

号外!号外!现在人们终于可以在浏览器中进行人脸识别了!本文将为大家介绍「face-api.js」,这是一个建立在「tensorflow.js」内核上的 java...

56120
来自专栏量子位

GitHub上大热的Deep Photo终于有TensorFlow版了!

王小新 编译整理 量子位 出品 | 公众号 QbitAI Prisma这个应用,你可能很熟悉。这是一个能将不同的绘画风格,迁移到照片中,形成不同艺术风格的图片。...

44260
来自专栏深度学习那些事儿

浅谈深度学习中超参数调整策略

深度学习中,设计模型以及保证模型的正确性是首要需要考虑的。当模型设置完成时,理论上模型不存在问题,实现效果也通过计算可以复现出来。一切准备就绪后,那么接下来需要...

23450
来自专栏CreateAMind

代码解读Top-down Neural Attention by Excitation Backprop及模型架构图

41330

扫码关注云+社区

领取腾讯云代金券