Facebook開源Mask R-CNN的PyTorch 1.0基準,比mmdetection更快、更省記憶體

機器之心發表於2018-10-25

專案地址:https://github.com/facebookresearch/maskrcnn-benchmark

Facebook開源Mask R-CNN的PyTorch 1.0基準,比mmdetection更快、更省記憶體

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 連結獲取。

Facebook開源Mask R-CNN的PyTorch 1.0基準,比mmdetection更快、更省記憶體

和 Detectron、mmdetection 的效能對比

訓練速度

下表中的資料單位是秒/迭代,越低越好。(mmdetection 中備註的硬體和 maskrcnn_benchmark 是不同的)

Facebook開源Mask R-CNN的PyTorch 1.0基準,比mmdetection更快、更省記憶體

訓練記憶體(越低越好)

Facebook開源Mask R-CNN的PyTorch 1.0基準,比mmdetection更快、更省記憶體

推理準確率(越高越好)

Facebook開源Mask R-CNN的PyTorch 1.0基準,比mmdetection更快、更省記憶體

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

相關文章