一塊GPU就能訓練語義分割網路,百度PaddlePaddle是如何優化的?

AI科技大本營發表於2019-01-31

一、 影象語義分割模型DeepLab v3

隨著計算機視覺的發展,語義分割成為了很多應用場景必不可少的一環。 比如網路直播有著實時剔除背景的要求,自動駕駛需要通過語義分割識別路面,與日俱增的應用場景對語義分割的精度和速度的要求不斷提高。同時,語義分割資料集也在不斷地進化,早期的Pascal VOC2,其解析度大多數在1000畫素以下。而Cityscape的語義分割資料集解析度全部達到了1024*2048,總共5000張圖片(精細標註),包含19類。這些資料集對研究者,計算裝置,甚至框架都帶來了更大的考驗。

DeepLab v3+ 是DeepLab語義分割系列網路的最新作,其前作有 DeepLab v1,v2, v3, 在最新作中,Liang-Chieh Chen等人通過encoder-decoder進行多尺度資訊的融合,同時保留了原來的空洞卷積和ASSP層, 其骨幹網路使用了Xception模型,提高了語義分割的健壯性和執行速率。其在Pascal VOC上達到了 89.0% 的mIoU,在Cityscape上也取得了 82.1%的好成績,下圖展示了DeepLab v3+的基本結構4:

一塊GPU就能訓練語義分割網路,百度PaddlePaddle是如何優化的?

DeepLab v3+在主幹網路之後連線了Encoder和Decoder,能夠在擴大網路感受的同時獲得更加高清的分割結果。

在PaddlePaddle的模型庫中已經包含了DeepLab v3+的訓練以及測試的程式碼。我們首先安裝最新版本的PaddlePaddle並且下載PaddlePaddle的模型庫:

pip install -U paddlepaddle git clone https://github.com/PaddlePaddle/models.git

當模型倉庫成功克隆,你將可以在目錄fluid/PaddleCV/deeplabv3+ 下看到用於訓練以及測試的程式碼:

ls fluid/PaddleCV/deeplabv3+ # ├── models.py # 網路結構定義指令碼 # ├── train.py # 訓練任務指令碼 # ├── eval.py # 評估指令碼 # └── reader.py # 定義通用的函式以及資料預處理指令碼

二、開始訓練

當資料和程式碼都已經準備好,我們可以開始訓練了,訓練的引數和指令如下:

CUDA_VISIBLE_DEVICES=0 FLAGS_fraction_of_gpu_memory_to_use=0.99 inplace_normalize=1 fuse_relu_before_depthwise_conv=1 python3 /home/cjld/nfs/liangdun/deeplabv3+/train.py --batch_size=-1 --train_crop_size=-1 --total_step=10 --base_lr=0.005 --train_set=train --norm_type=gn --save_weights_path=$YOUR_SAVE_WEIGHTS_PATH –dataset_path=$YOUR_DATASET_PATH

在這個命令中,我們沒有使用任何預訓練模型,從噪音開始訓練DeepLab v3+。並且是直接使用全解析度進行訓練(1024x2048,batch size=1)。幾個比較關鍵的引數解釋如下:

  • 環境變數CUDA_VISIBLE_DEVICES=0限制了訓練過程僅使用一張GPU,如果存在多張GPU,可以通過修改引數來得到訓練速度的提升。

  • 環境變數FLAGS_fraction_of_gpu_memory_to_use=0.99, 該環境變數將會讓PaddlePaddle佔用99%的視訊記憶體,可以根據實際情況進行調節。

  • 環境變數inplace_normalize=1,該引數是PaddlePaddle進行視訊記憶體優化的關鍵,開啟該開關將會讓框架對normalize layer進行 inplace 操作來優化視訊記憶體,現在支援的 normalize layer 有 group normalize。

  • 環境變數fuse_relu_before_depthwise_conv=1,該引數是視訊記憶體優化的另一個關鍵引數。這個引數會融合relu和depthwise conv來優化視訊記憶體。

  • 引數--save_weights_path=$YOUR_SAVE_WEIGHTS_PATH, 這裡你需要填入儲存模型的路徑。

  • 引數--dataset_path=$YOUR_DATASET_PATH, 這裡你需要填入資料集的路徑。

三、空間時間消耗分析

根據列印出來的資訊,我們可以發現,PaddlePaddle在訓練DeepLab v3+時,輸入一張全解析度的圖片,視訊記憶體消耗為10.2GB。得益於視訊記憶體消耗小於11G,我們可以使用1080ti完成訓練,訓練中每次迭代速度約為0.85s。

我們還可以使用工具,分析DeepLab v3+各部分視訊記憶體消耗情況:

一塊GPU就能訓練語義分割網路,百度PaddlePaddle是如何優化的?

在該圖中,視訊記憶體消耗最多的是主幹網路,佔用了68.1%,其次是decoder部分,佔用了16.4%,以及encoder佔用了5.3%,剩下其他部分為損失函式和資料預處理的視訊記憶體消耗。

視訊記憶體消耗最多的是主幹網路,佔用了68.1%,我們可以繼續檢視主幹網路內部視訊記憶體消耗情況:

一塊GPU就能訓練語義分割網路,百度PaddlePaddle是如何優化的?

這個圖展示了在主幹網路中的視訊記憶體消耗,Xception主幹網路主要由三部分組成, EntryFlow,MiddleFlow,以及ExitFlow,可以發現視訊記憶體消耗最大的是EntryFlow。

在上圖中,我們可以發現,儘管Xception主幹網路層數最多的部分是MiddleFlow,但是視訊記憶體消耗最大的卻是EntryFlow,這是因為在EntryFlow裡的特徵還沒有被充分下采用,解析度相當高,同時EntryFlow裡的通道數也不容小覷,因此造成了EntryFlow巨大的視訊記憶體開銷,這也為我們的優化指明瞭方向。

一塊GPU就能訓練語義分割網路,百度PaddlePaddle是如何優化的?

DeepLab v3+使用的主幹網路 Xception, MiddleFlow中的分離卷積塊重複了16次,層數相比較EntryFlow和ExitFlow要多得多,然而視訊記憶體消耗最大的卻是EntryFlow。

框架對比

除開對網路內部的視訊記憶體消耗進行分析,我們還對不同框架的視訊記憶體消耗進行了對比,下表展示了PaddlePaddle和TensorFlow1.12的視訊記憶體消耗以及效能對比,以下對比實驗使用的輸入資料是1024x2048全解析度的圖片,batch size為1,測試裝置P40(24G):

一塊GPU就能訓練語義分割網路,百度PaddlePaddle是如何優化的?

四、優化原理

這裡我們採用的視訊記憶體優化策略是 fuse_relu_before_depthwise_conv 和 inplace_normalize。顧名思義,fuse_relu_before_depthwise_conv 是講relu和depthwise_conv融合為同一個operator, 達到視訊記憶體的節省。而 inplace_normalize 則是使用原地操作來節省視訊記憶體。在卷積神經網路中,conv+normalize+activation是常見模式,在這種模式下,使用這兩種優化策略,可以節省3倍的視訊記憶體。這兩種優化策略歸納起來就是操作融合和原地計算, 是視訊記憶體優化中的常見策略, 對於不同的框架常常需要耗費人力進行開發, 而我們通過paddle的視訊記憶體優化可以很輕鬆的實現這一點。 該優化圖示如下:

一塊GPU就能訓練語義分割網路,百度PaddlePaddle是如何優化的?

在該圖中,紅框標註的data為會消耗視訊記憶體的資料塊,可以看到,通過inpalce和fuse兩種操作,原來需要儲存6個資料塊,優化後僅僅需要2個資料塊。

在上圖中,我們可以發現,inplace和fuse兩種操作,都分別幫助我們在每一個conv+normalize+activation塊中節省了1個資料塊,所以視訊記憶體節省的更多了,我們的分割網路也可以消耗更少的視訊記憶體資源。

總結

視訊記憶體空間優化和時間優化一樣,對空間消耗的壓榨是沒有止境的。層出不窮的模型,各種不同的優化方式,都對框架提出了很高的要求,選擇一款兼顧效果和效能的深度學習框架,往往能夠讓專案事半功倍。

1. The cityscapes dataset for semantic urban scene understanding.

2. The pascal visual object classes challenge.

3. The Mapillary Vistas Dataset for Semantic Understanding of Street Scenes.

4. Encoder-Decoder with Atrous Separable Convolution for Semantic Image Segmentation.

相關文章