作者 | Video++極鏈科技AI實驗室謝錦濱

整理 | 包包

個人對移動端神經(jīng)網(wǎng)絡(luò)開發(fā)一直饒有興致。去年騰訊開源了NCNN框架之后，一直都在關(guān)注。近期嘗試將分割網(wǎng)絡(luò)移植到NCNN，能夠在手機(jī)端實現(xiàn)一些有趣的應(yīng)用，本文就幾個技術(shù)話題作相關(guān)介紹。

神經(jīng)網(wǎng)絡(luò)選擇

目前做segmentation常用的一些神經(jīng)網(wǎng)絡(luò)網(wǎng)絡(luò)有如下幾個可以選擇:

• MaskRCNN

• FCN

• UNET

• SegNet

• Tiramisu

在移動端做人像分割有兩大優(yōu)勢，首先是隱私，其次是可以做到實時，能夠創(chuàng)造更多玩法。因為UNET模型比較簡單，干脆就從這個入手。下面是UNET網(wǎng)絡(luò)結(jié)構(gòu)：

首先我采用了基于keras的版本: https://github.com/TianzhongSong/Person-Segmentation-Keras，訓(xùn)了一個基本模型，大小為39M, iphone X上15秒處理一幀。明顯這個速度太慢，需要進(jìn)行改造。

移動端Inference框架

經(jīng)過調(diào)研，粗略比較了幾個神經(jīng)網(wǎng)絡(luò)框架:

其中使用難易程度，主要跟我個人習(xí)慣有關(guān)。NCNN框架比較好，代碼不多，而且兼容iOS和安卓（臺式機(jī)以及嵌入式環(huán)境同樣支持），同時底層計算采用匯編做了優(yōu)化。NCNN只實現(xiàn)神經(jīng)網(wǎng)絡(luò)的forward部分，沒有反向傳播，所以訓(xùn)練仍舊依賴其他開源框架，現(xiàn)在幾大框架都遵守ONNX協(xié)議，理論上各種框架模型之間互相轉(zhuǎn)換并不存在什么問題，工具也都是開源的。

不過keras沒辦法直接轉(zhuǎn)成ncnn模型，研究過通過onnx模型做中間跳板，采用了一些開源的轉(zhuǎn)換工具，也是一堆問題。NCNN支持幾個神經(jīng)網(wǎng)絡(luò)訓(xùn)練框架:caffe/mxnet/pytorch，在ncnn的github有一篇issue里nihui推薦采用MXNET，因此MXNET也成為了我的首選。其他框架往NCNN轉(zhuǎn)換工具:

NCNN轉(zhuǎn)換Tensorflow模型有問題； Caffe沒有Pytorch和MXNET好用; 最終在MXNet和Pytorch之間選擇了MXNet。

人像數(shù)據(jù)集

• https://github.com/lemondan/HumanParsing-Dataset

• https://github.com/ZhaoJ9014/Multi-Human-Parsing_MHP

• COCO人像數(shù)據(jù)集 – 加入后效果質(zhì)的飛躍

• ADE20K

網(wǎng)上找了上面幾個數(shù)據(jù)集，抽取出人像部分，采用基本的flip/crop/rotate操作做了擴(kuò)充，得到228423張訓(xùn)練樣本，另外湊了9064張驗證樣本。

模型轉(zhuǎn)換（MXNET->NCNN）

MXNET的UNET版本并沒有現(xiàn)成可用的合適版本。參照其他版本的UNET，自己coding完成一個版本。代碼請參考: https://github.com/xuduo35/unet_mxnet2ncnn。

在這個基礎(chǔ)上訓(xùn)練完成，用來測試ncnn轉(zhuǎn)換基本可用。這里提一下轉(zhuǎn)換過程遇到的一些問題和解決方案。

一個是調(diào)用ncnn extract函數(shù)會crash，經(jīng)過調(diào)查，發(fā)現(xiàn)mxnet2ncnn工具也有bug，blob個數(shù)算錯，其次是input層one_blob_only標(biāo)志我的理解應(yīng)該是false，不知道什么原因轉(zhuǎn)換過來的模型這邊是true，導(dǎo)致forward_layer函數(shù)里面bottoms變量訪問異常。后來一層層extract出來打印輸出的channel/width/height調(diào)查后又發(fā)現(xiàn)，自己代碼里unet.py里的name為pool5寫成了pool4，前面的crash跟這個致命錯誤有關(guān)系也有直接關(guān)聯(lián)。

第二個問題是轉(zhuǎn)成ncnn后的預(yù)測結(jié)果死活不對。只能一層層去檢查，寫了幾個簡單的工具可以打印中間隱藏層的結(jié)果（代碼: https://github.com/xuduo35/unet_mxnet2ncnn/check.py）。在這個基礎(chǔ)之上，發(fā)現(xiàn)是第一次反卷積就出了問題（mxnet神經(jīng)網(wǎng)絡(luò)trans_conv6的輸出）。結(jié)果完全不一致，按個人理解，反卷積算法會出問題的可能性基本為0，所以把mxnet這一層的權(quán)重值打印了出來。再在mxnet2ncnn的代碼里把對應(yīng)的參數(shù)打印，最后發(fā)現(xiàn)是num_group出了問題，簡單處理就是把mxnet2ncnn.cpp里的反卷積num_group固定為1，終于解決問題。得到正確的輸出結(jié)果：

中間還遇到一些ncnn和mxnet之間圖像格式之類的轉(zhuǎn)換問題，特別是浮點數(shù)的處理，就不啰嗦了。另外，調(diào)試過程發(fā)現(xiàn)，ncnn的中間層輸出和mxnet的輸出不是完全一致，可能是有一些參數(shù)或者運(yùn)算細(xì)節(jié)問題，不影響最后mask結(jié)果，暫時忽略。

幾個問題

到目前為止還存在幾個問題，1. 模型比較大；2. 單幀處理需要15秒左右的時間（Mac Pro筆記本，ncnn沒有使用openmp的情況）；3. 得到的mask結(jié)果不是特別理想。針對這三個問題，對網(wǎng)絡(luò)結(jié)構(gòu)進(jìn)行調(diào)整。

1. 模型比較大

采取將網(wǎng)絡(luò)卷積核數(shù)量減少4倍的方式，模型大小下降到2M，粗略用圖片測試，效果也還可以。同時把之前用0值填充圖片的方式，改成用邊界值填充，因為測試的時候發(fā)現(xiàn)之前的方式總在填充的邊界往往會出現(xiàn)檢測錯誤。

2. 單幀處理需要15秒左右的時間

按照第一步處理之后，基本上一張圖片只要1秒鐘就處理完成，。在手機(jī)上由于NCNN做了優(yōu)化，經(jīng)過測試速度是Mac Pro的好幾倍。

3. 得到的mask結(jié)果不是特別理想

在權(quán)衡模型大小和準(zhǔn)確率的基礎(chǔ)上修改UNET網(wǎng)絡(luò)結(jié)構(gòu)，具體不再贅述。

最終結(jié)果