自動網路搜尋(NAS)在語義分割上的應用(二)
阿新 • • 發佈:2020-05-26
前言:
本文將介紹如何基於ProxylessNAS搜尋semantic segmentation模型,最終搜尋得到的模型結構可在CPU上達到36 fps的測試結果,展示自動網路搜尋(NAS)在語義分割上的應用。 隨著自動網路搜尋(Neural Architecture Search)技術的問世,深度學習已慢慢發展到自動化設計網路結構以及超引數配置的階段。尤其在AI落地的背景下,許多模型需要部署在移動端裝置。依據不同裝置(GPU, CPU,晶片等),不同的模型需求(latency, 模型大小,FLOPs),使用NAS自動搜尋最佳網路結構將會是一個很有潛力的方向。上一篇介紹了NAS的基本框架和入門必讀DARTS [1],以及在semantic segmentation領域的應用。 距離現在才不過幾個月,NAS論文數量明顯增長:在理論研究方面,search strategy,evaluation performance的方法看似趨於穩定,不得不提到最近FAIR團隊的RegNet [2]探討了搜尋空間的設計,通過大量實驗把常見的設計模型的理論一一驗證,我們可以根據它的結論縮小搜尋空間從而提高搜尋效率;在應用方面以還是object detection為主,也有segmentation, reID, GAN等領域。 NAS算是一項新技術,但是語義分割semantic segmentation卻是老生常談。自FCN的問世起,SegNet, UNet這種簡單粗暴的encoder-decoder結構在多種影象上都能達到可以接受的結果,deeplab系列之後更是在開源資料集達到巔峰。從學術角度看semantic segmentation似乎已達到瓶頸,於是researcher們紛紛轉向小樣本,semi-supervised,domain adaption, cloud point等方向另闢蹊徑。但是semantic segmentation落地卻非常困難。在實際落地場景中,使用常見的backbone (resnet或yolo系列) 就能夠完成各種object detection任務,但是在segmentation上效果卻不好:- 由於光線等原因,實際場景影象的intensity分佈更復雜,而segmentation需要細分邊界, 對畫素值的判定尤為重要。然而,相比detection來說segmentation的資料標記成本高導致訓練資料較少,只依靠data augmentation等手段提升有限。
- Segmentation是pixelwise的任務,因為它要處理到每一個pixel,所以模型一般都會比object detection的模型大許多(你看這個模型它又長又寬)。如果你的模型被要求real-time推理(>16 fps),那麼準確度和速度必然會成為衝突,Double kill!
- 當語義分割用在了視訊流,對準確度的要求會更高。即使每兩幀只相差幾個pixel,即使在mIoU的數值上相差無幾,但是人眼看上去不夠穩定,會有“抖動”的邊界, Triple kill!
- 當語義分割模型走下雲端,部署在算力有限的移動端,底層晶片可能對很多操作不支援,使得在原本在可以在GPU上開心玩耍的模型到了CPU上便一朝打回解放前, Quadra kill!
1.Overview of ProxylessNAS
- Search Space: 在搜尋空間中定義的operation candidate的是來自MobileNetv2 [8]的block,分別取不同的kernel size(3, 5, 7)和不同的expansion rate(3,6),再加上identity和zero操作一共8種ops(c.f. Figure 1)。 網路的巨集觀結構是一個常見的鏈狀結構來完成classification, 每一層都有8個ops candidate(c.f. Figure 2)。正如前面提到的,運算元之間太複雜的連線方式會讓速度變慢,常見的小模型結構都是這種鏈狀結構。
- Search Strategy: 搜尋策略採用可微分的方法,這種搜尋策略近兩年很常見。雖然不及RL和EA穩定,但是可以大幅度提高搜尋速度。
- Evaluation Performance: One-shot 權值共享, 也是現有最常見的super-net的形式。對於計算資源匱乏的團隊和個人來說,這種方式能夠提高搜尋效率減少記憶體佔用。
- Training: 每個iteration開始的時候,在每一層都隨機啟用一個operation(c.f. the binary gate in Figure 2),將所有啟用的operation連線起來組成一條子網路記為subnet,通過back propagation來更新這條subnet的weight。沒有啟用的ops不放入記憶體,也就是說訓練的時候只有一條subnet在記憶體中,這也使得整個搜尋過程可以在單卡上完成。
- Searching:每個operation的權重alpha代表它的重要程度,也就是最終被選擇的概率,probability = Softmax(alphas)。換言之,搜尋的過程,就是不斷更新權重alpha的過程。和training一樣,每個iteration都要隨機啟用一條subnet,但是這次要讓operation的weight固定,通過back propagation計算這條subnet上的alpha。Paper裡面Eq (4)給出了計算方式,由於binary gate和probability成正比,公式裡面將loss對probability的求導轉化成對binary gate的求導,而loss對binary gate的導數在back propagation的時候有計算過並且儲存了下來(這部分paper沒有細說可參考原始碼)。
2.Real-time Semantic Segmentation using ProxylessNAS on CPU
儘管對ProxylessNAS還有很多沒有解決的問題,奈何單卡搜尋訓練省時省力瑕不掩瑜。藉助Intel的openvino推理框架,本文嘗試用ProxylessNAS搜尋可執行在CPU(x86)上的real-time semantic segmentation模型做人形分割,下面會詳細介紹對演算法的改進和實驗結果。 2.1 Super-net setting- Search space: 在設定搜尋空間的時候,秉著大力出奇跡的心態我把常用的operation都塞了進來,分別是MBv3 (3x3), MBv3 (5x5), DilConv (3x3), DilConv (5x5), SepConv (3x3), SepConv (5x5), ShuffleBlock一共7種ops。其中MBv3是來自MobileNetv3 [5]的基本模組,DilConv和SepConv是來自DARTS [1]的dilated sepatable convolutions和separable convolutions,ShuffleBlock是來自ShuffleNetv2 [4]的基本模組,前面三種operation都設定了兩種kernel size可以選擇。在定義巨集觀網路結構的時候,採用deeplabv3+ [9]的結構 (c.f. Figure 3): head + searchable backbone + ASPP + decoder。與UNet類似,將encoder的feature map直接”add”到decoder,這裡沒有用”concatenation”是為了避免模型過“寬”使速度變慢。其中s2, s4, s8, s16, s32分別指feature map的resolution下降2,4,8,16,32倍。與ProxylessNAS類似,supernet的引數包含兩部分,一部分是operation本身的weight,另一部分是operation的權重alpha。
- Searching Strategy: 延續ProxylessNAS的可微分求導方式
- Evaluation Performance: One-shot權重共享
- Decoupling the training and searching process: 在ProxylessNAS中“training”和“searching”是同時輪流完成的,也就是一邊訓練一邊搜尋。我在實驗的時候把“training”和“searching”徹底分開,先用50個epochs只更新super-net裡面operation的引數,在訓練之後,再更新operation的權重alphas。這麼做的原因是避免在operation引數不穩定的時候,某些alpha過大影響後面的決策。
- Consider the latency as a hard constraint: 因為模型推理速度比較重要,而且不能用簡單的疊加方式計算,所以每次隨機啟用subnet的時候都要算一下這條subnet的推理速度,如果不符合要求(如latency > 30ms)則重新搜尋一條subnet,這樣一定程度上避免很多推理速度過慢的operation被選擇和學習。
- Task: 基於CPU(x86)的實時人像分割
- DL platform: Intel openvinohttps://software.intel.com/content/www/us/en/develop/tools/openvino-toolkit.html
- Dataset : >20k張影象,一部分來自 coco/pascal資料集中帶有”person”類別的,另一部分是私有資料
- Data augmentation: random crop, cutout, random brightness/contrast adjust, random Gaussian blur/sharpen
- Searching time: 單卡2 GPU days (K80) 包括training和searching
3.Future work
通過實驗我們看到ProxylessNAS搜尋策略可以從classification遷移到segmentation,在速度相仿的情況下,搜尋出來的網路要比原本MobileNetv3 [5]準確度提高很多。但是隻限於當前的場景,不能說人工設計出來的模型就不好或一定會被取代(雖然MobileNetv3也是NAS搜出來的)。在特定場景和有特定需求的時候,用NAS設計網路結構確實比人工設計加上大量調參實驗要更高效,在AI落地方面更有發展前景。本文只是初探ProxylessNAS,後續還會有以下幾個方面的探索。- 實驗結果表明super-net權值共享的形式有一定合理性。但是在結構搜尋的時候,將每層probability最大的operation組成subnet作為輸出結果還是有不合理之處。因為subnet在搜尋和訓練的時候具有一定的耦合性,每層的operation一榮俱榮一損俱損。最終將每層最佳的operation選出來,組合在一起的時候未必能符合預先設定的hard constraint,這裡還是有需要改進的地方,比如可以計算相鄰兩層operation的sub-path的權重代替每層operation的權重。
- ProxylessNAS是MIT Hansong團隊早期的work,現在已有後續OFA問世(也是跪著讀完的)。在OFA中作者徹底將training和searching分開,結合了knowledge distillation,先訓練teacher model,然後用NAS的思路在teacher model中搜索出最佳student model。OFA可以理解為自動化network pruning或自動distillation。如果OFA實驗效果好,後續還會有關於OFA的實戰經驗的分享。
- Figure 5種實際效果展示的時候,人像和背景融合的比較自然,但是語義分割歸根到底是一個分類任務,邊緣的pixel“非黑即白”,如果想要和背景自然的融合,需要計算出前景的透明度alpha matte,這裡涉及到另一項背景摳圖技術,和segmentation配合使用效果更佳。其實Figure 5的下圖中已經看出segmentation沒有把頭髮分割出來,但是在結果中卻保留了下來,也是用了背景摳圖的原因。Matting除了可以優化segmentation結果,還可以實現切換背景(cf. Figure 7),PS等功能。