基于NVIDIA Jetson TX2的道路場景分割①

李詩菁， 卿粼波， 何小海， 韓 杰

(四川大學 電子信息學院， 成都 610065)

近年， 深度學習的迅猛崛起給各個科技領域的發展帶來了巨大影響. 在計算機視覺領域， 圖像語義分割是重要的基礎研究問題之一. 傳統的圖像分割一般采用基于閾值、邊緣、區域的分割方法， 根據顏色紋理等人工標定的特征完成分割， 過程復雜且局限性較大[1]. 基于深度學習的圖像語義分割網絡的分割性能遠超傳統算法， 可應用到諸多領域如無人駕駛、智能安防、交通監控、機器人中[2]. 在較有代表性的圖像分割網絡FCN (Fully Convolutional Networks)[3]、SegNet(SegNet: A Deep ConvolutionalEncoder-Decoder Architecture for ImageSegmentation)[4]、PSPNet(Pyramid Scene Parsing Network)[5]和ENet(ENet: A Deep Neural Network Architecture forReal-Time Semantic Segmentation)中， FCN是最早的圖像分割網絡之一， 準確度較低且耗時長； PSPNet分割的像素精確度高， 但耗時較長； ENet分割速度較快， 但分割效果較不理想； SegNet兼具較高的分割準確度與速度， 是完成道路場景分割任務較好的選擇.

隨著人工智能的不斷發展， 深度學習網絡在嵌入式平臺上部署的需求也日益增長. 目前國內外推出的可部署深度學習的主流芯片一般分為FPGA芯片和GPU芯片兩類. 在FPGA上已實現了一些小規模的神經網絡， 如Zhang等人在Xilinx Virtex 7485T FPGA上實現了AlexNet[6]， 盧冶、陳瑤等人在XC7Z020 FPGA上實現了LeNet-5、Cifarnet網絡[7]. 基于FPGA的深度學習應用的功耗小于GPU， 但受限于FPGA開發難度與硬件資源， 目前大規模復雜結構的網絡在嵌入式上的部署還是更多地基于GPU實現. 已有許多深度學習網絡逐漸被應用到了嵌入式ARM+GPU平臺如NVIDIA JetsonTX2 (簡稱TX2)上， 構建了目標檢測、感知導航等系統[8，9]. 因分割網絡相較于圖像分類、目標檢測更難實現， 大部分高準確度的分割網絡在低功耗的嵌入GPU平臺上運行速度極慢， 而較少地被應用到嵌入式平臺上.

因此， 針對無人駕駛與交通監控等應用背景， 本文選取SegNet網絡在嵌入式平臺TX2上實現道路場景理解， 并采用模型簡化與轉換網絡模型為加速引擎的方式， 在基本保持原網絡分割準確度的情況下大幅度提升了分割網絡在嵌入式平臺中的推理速度. 本文首先合并BN層簡化了網絡模型， 然后基于英偉達的TensorRT2前向推理加速器采用水平層集成， 關聯層消除， 權值精度減半等優化措施， 將網絡模型轉換成了加速引擎進行前向推理. 下面將分別就網絡結構與模型簡化、基于TensorRT2的加速引擎構建兩部分進行介紹.

1 網絡結構與網絡模型簡化

1.1 網絡結構

本文采用SegNet網絡， 實現將道路場景分為行人、道路、天空、標志、車輛、交通標線、桿狀物、自行車、側路、植物、建筑、圍墻共12類的任務.

SegNet網絡基本結構為自動編-解碼器結構， 采用VGG16前13層卷積層作為編碼器， 后接13層解碼網絡與一個分類層. SegNet關鍵在于存儲了編碼網絡中每個的池化層中的最大值與其空間信息， 用于對應解碼器的非線性上采樣， 這極大精確了分割中的邊界定位， 減少了編碼器到解碼器的參數量， 使得SegNet在速度與內存利用上都具有很大優勢.

1.2 網絡模型簡化

在將網絡模型轉換為加速引擎之前， 本文對SegNet網絡進行了模型簡化. SegNet網絡中引入了BN層， 這種層可在訓練時期加速網絡收斂， 但在推理過程中會增大內存消耗， 降低推理速度. 可將BN層與其相連的卷積層合并以簡化模型加速推理， 如圖1所示.

圖1 BN層合并

卷積層與BN層都是線性變換， 提供了合并的可能性. 設輸入的mini-batch: {x1··n}， BN層變化如式(1):

式(1)中，i∈[0,n]， 且i為整數與是訓練中得到的參數為趨于0的常數.

2 基于TensorRT2的加速引擎構建

SegNet在TX2上的推理速度極慢， 因此本文采用NVIDIA推出的TensorRT2對網絡模型進行加速.

2.1 基本流程

TensorRT是NVIDIA公司推出的深度學習網絡推理引擎， 可優化已有網絡模型， 大幅提升神經網絡在如機器人、自動駕駛平臺上的推理速度. 目前NVIDIA公司已推出四個版本的TensorRT. 不同版本TensorRT可支持的深度學習框架不同.

因為本文使用的SegNet網絡采用caffe架構， 且需要插入TensorRT不支持的層， TensorRT2版本已經能滿足本文需求. 故本文使用TensorRT2， 基于訓練后的網絡模型生成加速推理引擎， 完成高效的GPU推理性能優化. 加速后的模型無需深度學習架構支持， 對平臺的依賴性極小， 無論在高性能GPU(如NVIDIA Titan X)上， 還是在嵌入式平臺(例如NVIDIA TX2)上都擁有不俗的加速能力. 加速引擎構建的基本流程如圖2.

圖2 加速引擎構建流程

構建過程中首先創建網絡， 導入網絡模型并利用NvcaffeParser進行解析， 然后根據網絡結構與定義采取優化措施生成加速引擎. 生成的加速引擎可序列化存儲到磁盤中. 在執行階段反序列化加速引擎， 進行輸入與輸出口的綁定、GPU顯存分配與計算內核啟動，執行加速引擎得到分割結果.

2.2 自定義層添加

在上文的網絡模型解析過程中， 需將原網絡模型中的層轉換為其支持的形式. 加速引擎中除部分常見層如卷積層、激活層、全連接層等可直接轉換外， 其它TensorRT不支持的層需自行定義， 下文將這些層統稱為自定義層. TensorRT中支持的層可直接由上文提到的NvcaffeParser進行解析， 自定義層則需采用plugin接口添加入NvcaffeParser中.

自定義層plugin添加的主要流程如下: 輸出確定、層配置、工作空間分配、資源管理、序列化與層執行. 在加速引擎中， 一個層的輸入或輸出定義為tensor， tensor具有數據類型與三個維度分量， 即通道數C、寬度W與高度H. 在開始插入層時， 需對插入層的輸出進行定義， 需確定插入層的輸出數目與輸出tensor的三個維度分量{C，H，W}； 完成輸出確定后， 將進行層配置， 該過程主要獲取輸入tensor的形式； 在層配置后， 加速引擎會分配臨時的工作空間在自定義層之間共享以達到內存利用率最大化； 同時還需要進行資源管理的配置， 主要是通過層資源初始化與銷毀來完成資源分配與釋放； 在進行完成上述部分后， 需判斷加速引擎是否需要序列化到磁盤中， 若需序列化， 則將自定義層的參數與網絡的其余部分合并以便后續整個網絡的序列化存儲； 最后在層執行階段， 主要完成層的算法實現， 如果未選擇序列化存儲， 則直接在資源分配后執行該過程， 若已進行序列化存儲， 則提取序列化參數后執行該過程， 整體過程如圖3所示.

圖3 plugin層插入流程

本文將網絡中需兩個輸出口的pooling層、upsample層與argmax分類層添加到加速模型中， 這些添加的層是網絡的關鍵結構， 如圖4所示. 本文將pooling層分出了兩個輸出口， 將池化結果送入下一級編碼器的同時將池化的最大特征值與位置信息輸入了對應upsample層中. upsample層將接收的池化最大特征值按照位置信息填入上采樣的稀疏特征圖中， 空缺位置補零， 作為解碼器的輸入， 如圖5. 最終經過層層解碼后的特征圖輸入argmax層， 得到每個像素最大可能性的分類結果.

圖4 plugin層結構

圖5 plugin層實現過程

2.3 優化措施

本文的優化措施主要分為圖1中的引擎生成階段的優化與圖3自定義層添加中執行階段enqueue( )時的CUDA編程優化.

網絡模型中的各個層經過解析之后， 在引擎生成時主要采用了以下措施進行優化: 清除未使用的輸出層以避免不必要的計算； 整合垂直結構的卷積層、ReLU激活層、Bias偏置項為CBR層； 將擁有同一輸入、相同操作、相似參數和不同輸出的水平層集成為一個層，例如圖6； 在平臺支持的情況下， 將數據精度需求從32位float降低為半精度的fp16， 以此提升計算效率.

在自定義層添加時執行層算法的函數enqueue( )中， 本文采用了CUDA編程的方式將算法中大量耗時的計算分配到GPU上并行實現. 原網絡最后的分類層argmax層為CPU實現， 若同樣采用CPU實現， 需將輸入層的數據從GPU上復制到CPU上. 數據復制將耗費大量時間， 且argmax層算法在CPU上的執行時間也過長. 因此本文在GPU上實現了argmax層前饋算法. 在CUDA C中的host函數只可由CPU調用， device函數則只可運行在GPU上， 且不能調用常用函數. 因此argmax層前饋算法中調用的部分算法， 并不能在GPU上調用. 為此本文針對網絡的實際應用簡化了相應算法以便GPU實現. 最終在得到每個像素值12類分類結果后， 在一維索引空間中并行360*480個線程，將位置一一對應的特征值的12個分類結果進行排序，得到整幅圖像像素級別的分類.

圖6 水平層集成

3 實驗測試

3.1 網絡訓練

本文采用GPU模式進行網絡訓練， 硬件為基于Pascal架構的NVIDIA Titan X顯卡， 該卡顯存為12 GB.

為滿足針對交通監控與自動駕駛的道路場景理解的目的， 本文從包含城市交通場景的數據庫KITTI[10]、Camvid[11]與Cityscapes[12]中選取了1300張圖像作為訓練數據集， 并將原圖與標簽裁剪為統一尺寸. 此外， 選取的網絡訓練需要單通道標簽， 故將原數據集中RGB標簽轉換為單通道， 并將每個像素轉換為代表天空、行人、圍墻、植物、交通標志、側道、道路、車輛等12個分類的灰度值. 在訓練時， 為提高分類準確度， 使用預先訓練好的VGG16網絡模型的權值對SegNet編碼網絡的權值進行初始化. 在20 000次iteration后， 網絡準確度就已經達到80%以上， 在迭代36 000次時，達最高準確度92%， 迭代38 000次時達到最低損失0.019.

3.2 測試結果

在訓練得到最優網絡模型的基礎上， 本文完成了加速引擎(下文稱為SegEng)的構建， 并將其序列化后存儲到磁盤上. NVIDIA TX2平臺支持fp16， 故本文生成SegEng時采用了半精度模式， 將其大小縮減為SegNet模型的一半. 在 NVIDIA TX2 上， 本文對 GPU模式下使用cudnn6.0加速的SegNet與SegEng的前向推理性能進行了測試對比.

圖7 分割結果

分割圖對比如圖7所示. 在圖像分割領域， 常用以下三個指標衡量分割精度: 像素精確度(G)表示圖像中被正確分類的像素比； 類平均精確度(C)表示所有類被正確分類的平均值； 交除并均值(Mean Intersection over Union， MIoU)代表所有類的IoU平均值.

本文從測試集中選取圖片， 使用上述三個指標對迭代了40 000次得到的SegNet模型與SegEng的分割結果進行評估， 得到表1.

表1 分割精度對比(單位: %)

原始模型中的權值精度為32位， 加速引擎舍棄了權值中后16位的冗余信息， 將權值精度降低為16位，權值中丟棄的后16位信息對最終的像素級別分類概率影響極小， 由表1中可看到并未影響到最終的分割準確度.

在NVIDIA TX2上， 本文對不同分辨率的1000幀視頻進行了10次測試后取平均值， 得到了表2中SegNet與SegEng推理時間的對比.

表2 單幀圖像推理時間對比(單位: ms)

在經過上文的層合并、網絡層優化、CUDA加速與TensorRT內部的卷積優化、內存優化、精度減半等優化措施后， 由表 2可看出， 在 TX2上， 相較于SegNet， 本文生成的SegEng達到了約十倍的加速比，大幅度提高了SegNet在嵌入式平臺上的推理速度. 在分割 480*320大小的圖片時， 分割速度可達 10 fps， 具有在實際道路場景解析的應用潛力.

4 結語

本文針對無人駕駛與交通監控的應用背景， 采用了圖像語義分割網絡SegNet完成將道路場景分割為行人、車輛、道路、植物、側道、交通標志、自行車等12類對象的任務. 同時本文為提升分割網絡在嵌入式平臺上的推理速度， 對網絡模型進行了簡化， 然后基于NVIDIA推出的TensorRT2， 采用集成水平層、消除多余輸出層、采用權值精度減半、CUDA并行編程等優化措施， 完成了網絡模型加速并生成了加速引擎.在NVIDIA TX2嵌入式平臺上， 本文生成的加速引擎無需深度學習架構支持， 在分割精度基本無影響情況下， 推理速度可達原網絡的十倍， 為復雜結構分割網絡在嵌入式平臺上的應用提供了支持.