PyTorch中的MIT ADE20K數據集的語義分割
代碼地址:https://github.com/CSAILVision/semantic-segmentation-pytorch
Semantic Understanding of Scenes through ADE20K Dataset. B. Zhou, H. Zhao, X. Puig, T. Xiao, S. Fidler, A. Barriuso and A. Torralba. International Journal on Computer Vision (IJCV), 2018. (https://arxiv.org/pdf/1608.05442.pdf)
Scene Parsing through ADE20K Dataset. B. Zhou, H. Zhao, X. Puig, S. Fidler, A. Barriuso and A. Torralba. Computer Vision and Pattern Recognition (CVPR), 2017. (http://people.csail.mit.edu/bzhou/publication/scene-parse-camera-ready.pdf)
這是MIT ADE20K場景解析數據集(http://sceneparsing.csail.mit.edu/)上語義分割模型的PyTorch實現。
ADE20K是MIT Computer Vision團隊發布的最大的用於語義分割和場景解析的開源數據集。請通過下面的鏈接找到數據集和在Caffe和Torch7上的實現的存儲庫:https : //github.com/CSAILVision/sceneparsing
如果只是想玩演示,請嘗試以下鏈接:http : //scenesegmentation.csail.mit.edu可以上傳自己的照片並進行解析!
所有預訓練的模型都可以在以下位置找到:http ://sceneparsing.csail.mit.edu/model/pytorch
[從左至右:測試圖像,地面真理,預測結果]
可以在以下位置找到語義類別的顏色編碼:https ://docs.google.com/spreadsheets/d/1se8YEtb2detS7OuPE86fXGyD269pMycAWe2mtKUj2W8/edit?usp=sharing
更新
- 現在支持HRNet模型。
- 使用配置文件來存儲參數解析器中的大多數選項。選項的定義在中詳細介紹config/defaults.py。
- 在數據預處理中遵循Pytorch慣例(RGB [0,1],減去均值,除法std)。
強調
PyTorch上的同步批處理規范化
此模塊在訓練期間計算所有設備上的均值和標准差。憑經驗發現,合理的大批量批量對於細分很重要。感謝毛嘉元先生的貢獻,有關詳細信息,請參閱Synchronized-BatchNorm-PyTorch。
該實現易於使用,因為:
- 是純Python,沒有C ++額外的擴展庫。
- 與PyTorch的實現完全兼容。具體來說,使用無偏方差更新移動平均值,並使用sqrt(max(var,eps))代替sqrt(var + eps)。
- 非常有效,僅比UnsyncBN慢20%到30%。
動態輸入比例,可使用多個GPU進行訓練
對於語義分割的任務,最好在訓練過程中保持圖像的長寬比。因此,重新實現了該DataParallel模塊,並使其支持在python dict中將數據分發到多個GPU,以便每個gpu可以處理不同大小的圖像。同時,數據加載器的操作也不同。
現在,數據加載器的批處理大小始終等於GPU的數量,每個元素將被發送到GPU。也與多處理兼容。請注意,用於多處理數據加載器的文件索引存儲在主進程中,這與目標每個工作人員維護自己的文件列表的目標相矛盾。因此,使用了一個技巧,盡管主進程仍為數據加載器提供__getitem__功能索引,但只是忽略了此類請求並發送了一個隨機批處理指令。另外,如果直接使用上述技巧,則由數據加載器派生的多個工作程序都具有相同的種子,會發現多個工作程序將產生完全相同的數據。因此,添加了一行代碼,該代碼為numpy.random 在激活數據加載器中的多個工作程序之前。
最先進的模型
- PSPNet是場景解析網絡,通過金字塔池模塊(PPM)聚合全局表示。是ILSVRC'16 MIT場景解析挑戰賽的獲勝者模型。有關詳細信息,請參閱https://arxiv.org/abs/1612.01105。
- UPerNet是基於特征金字塔網絡(FPN)和金字塔池模塊(PPM)的模型。不需要膨脹卷積,運算符既耗時又消耗內存。無需花哨,與PSPNet相比甚至更好,同時需要更短的訓練時間和更少的GPU內存。有關詳細信息,請參閱https://arxiv.org/abs/1807.10221。
- HRNet是最近提出的模型,該模型保留了整個模型的高分辨率表示,而沒有傳統的瓶頸設計。在一系列像素標記任務上實現了SOTA性能。有關詳細信息,請參閱https://arxiv.org/abs/1904.04514。
支持機型
將模型分為編碼器和解碼器,其中通常直接從分類網絡修改編碼器,並且解碼器由最終卷積和上采樣組成。在config文件夾中提供了一些預配置的模型。
編碼器:
- MobileNetV2dilated
- ResNet18 / ResNet18dilated
- ResNet50 / ResNet50dilated
- ResNet101 / ResNet101dilated
- HRNetV2(W48)
解碼器:
- C1(一個卷積模塊)
- C1_deepsup(C1 +深度監控技巧)
- PPM(金字塔合並模塊,有關詳細信息,請參閱PSPNet。)
- PPM_deepsup(PPM +深度監督技巧)
- UPerNet(金字塔池+ FPN頭,有關詳細信息,請參閱UperNet。)
性能:
重要說明:存儲庫中的基礎ResNet是自定義的(不同於Torchvision中的ResNet)。基本模型將在需要時自動下載。
系統結構 |
多尺度測試 |
平均IOU |
像素精度(%) |
總分 |
推理速度(fps) |
MobileNetV2dilated + C1_deepsup |
沒有 |
34.84 |
75.75 |
54.07 |
17.2 |
是 |
33.84 |
76.80 |
55.32 |
10.3 |
|
MobileNetV2dilated + PPM_deepsup |
沒有 |
35.76 |
77.77 |
56.27 |
14.9 |
是 |
36.28 |
78.26 |
57.27 |
6.7 |
|
ResNet18dilated + C1_deepsup |
沒有 |
33.82 |
76.05 |
54.94 |
13.9 |
是 |
35.34 |
77.41 |
56.38 |
5.8 |
|
ResNet18dilated + PPM_deepsup |
沒有 |
38.00 |
78.64 |
58.32 |
11.7 |
是 |
38.81 |
79.29 |
59.05 |
4.2 |
|
ResNet50dilated + PPM_deepsup |
沒有 |
41.26 |
79.73 |
60.50 |
8.3 |
是 |
42.14 |
80.13 |
61.14 |
2.6 |
|
ResNet101dilated + PPM_deepsup |
沒有 |
42.19 |
80.59 |
61.39 |
6.8 |
是 |
42.53 |
80.91 |
61.72 |
2.0 |
|
UperNet50 |
沒有 |
40.44 |
79.80 |
60.12 |
8.4 |
是 |
41.55 |
80.23 |
60.89 |
2.9 |
|
UperNet101 |
沒有 |
42.00 |
80.79 |
61.40 |
7.8 |
是 |
42.66 |
81.01 |
61.84 |
2.3 |
|
HRNetV2 |
沒有 |
42.03 |
80.77 |
61.40 |
5.8 |
是 |
43.20 |
81.47 |
62.34 |
1.9 |
該訓練是在具有8個NVIDIA Pascal Titan Xp GPU(12GB GPU內存)的服務器上進行基准測試的,推理速度是對單個NVIDIA Pascal Titan Xp GPU進行可視化的基准測試。
環境
該代碼是在以下配置下開發的。
- 硬件:> = 4個GPU用於訓練,> = 1個GPU用於測試(進行相應設置[--gpus GPUS])
- 軟體:Ubuntu 16.04.3 LTS,CUDA> = 8.0,Python> = 3.5,PyTorch> = 0.4.0
- 依賴項:numpy,scipy,opencv,yacs,tqdm
快速入門:使用訓練有素的模型對圖像進行測試
- 這是一個對單個圖像進行推斷的簡單演示:
chmod + x demo_test.sh
./demo_test.sh
該腳本下載經過訓練的模型(ResNet50dilated + PPM_deepsup)和測試圖像,運行測試腳本,並將預測的分段(.png)保存到工作目錄。
- 要在圖像或圖像文件夾($PATH_IMG)上進行測試,只需執行以下操作:
python3 -u test.py --imgs $PATH_IMG --gpu $GPU --cfg $CFG
訓練
- 下載ADE20K場景解析數據集:
chmod + x下載_ADE20K.sh
./download_ADE20K.sh
- 通過選擇要使用的GPU($GPUS)和配置文件($CFG)來訓練模型。訓練期間,默認情況下,檢查點保存在文件夾中ckpt。
python3 train.py --gpus $ GPUS --cfg $ CFG
- 要選擇要使用的GPU,可以執行--gpus 0-7或--gpus 0,2,4,6。
例如,可以從提供的配置開始:
- 訓練MobileNetV2dilated + C1_deepsup
python3 train.py --gpus GPUS --cfg config / ade20k-mobilenetv2dilated-c1_deepsup.yaml
- 訓練ResNet50dilated + PPM_deepsup
python3 train.py --gpus GPUS --cfg config / ade20k-resnet50dilated-ppm_deepsup.yaml
- 訓練UPerNet101
python3 train.py --gpus GPUS --cfg config / ade20k-resnet101-upernet.yaml
- 也可以在命令行中覆蓋選項,例如 python3 train.py TRAIN.num_epoch 10 。
評價
- 在驗證集上評估訓練有素的模型。將VAL.visualize True參數添加到輸出可視化效果中,如預告片所示。
例如:
- 評估MobileNetV2dilated + C1_deepsup
python3 eval_multipro.py --gpus GPUS --cfg config / ade20k-mobilenetv2dilated-c1_deepsup.yaml
- 評估ResNet50dilated + PPM_deepsup
python3 eval_multipro.py --gpus GPUS --cfg config / ade20k-resnet50dilated-ppm_deepsup.yaml
- 評估UPerNet101
python3 eval_multipro.py --gpus GPUS --cfg config / ade20k-resnet101-upernet.yaml
與其項目整合
可以通過安裝該庫pip以輕松與另一個代碼庫集成
點安裝git + https:
pip install git+https://github.com/CSAILVision/semantic-segmentation-pytorch.git@master
現在,可以輕松地以編程方式使用此庫。例如
from mit_semseg.config import cfg
from mit_semseg.dataset import TestDataset
from mit_semseg.models import ModelBuilder,