當前位置:
首頁 > 新聞 > 基於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 的結果。


基於PyTorch的目標檢測工具箱,商湯聯合港中文開源mmdetection

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。

基於PyTorch的目標檢測工具箱,商湯聯合港中文開源mmdetection

2. 訓練速度

訓練速度的單位是 s/iter,數值越低代表速度越高。

基於PyTorch的目標檢測工具箱,商湯聯合港中文開源mmdetection

3. 推斷測試

推斷速度的單位是 fps (img/s),數值越高代表效果越好。

基於PyTorch的目標檢測工具箱,商湯聯合港中文開源mmdetection

測試與訓練

首先安裝 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


喜歡這篇文章嗎?立刻分享出去讓更多人知道吧!

本站內容充實豐富,博大精深,小編精選每日熱門資訊,隨時更新,點擊「搶先收到最新資訊」瀏覽吧!


請您繼續閱讀更多來自 機器之心 的精彩文章:

Ian Goodfellow專訪:我為什麼可以在一夜之間創建GAN
聯合學習離散句法結構和連續詞表徵的無監督方法

TAG:機器之心 |