基于優化神經網絡的遙感影像語義分割算法

王曉鋒 張猷 夏嚴峰 王森 王輝

摘要：利用機器學習技術可以自動化地識別地表覆蓋物，減輕人力成本和減少人為誤差，加快測繪、國防、農業和防災減災等各個領域的數字化轉型。遙感影像地表覆蓋物的分類算法已經得到了一定的發展，但具體到應用場景中，受制于模型效率，自動化水平仍不高。為了保證訓練產生的機器學習模型可以穩定、快速地服務于上述的領域，本文針對機器學習模型在訓練、測試階段的優化進行了探索，從而保證模型可更高效穩定地在有限的硬件資源上運行。通過本文介紹的優化方法，模型的訓練速度得到了提升，模型占用的空間得到了減小，模型的推理速度得到了提高，并且保證了模型的精度幾乎不受影響。

關鍵詞：神經網絡壓縮;神經網絡優化;深度神經網絡;圖像語義分割;遙感影像

中圖分類號：TP391? ? ? 文獻標識碼：A

文章編號：1009-3044（2022）17-0070-03

1 概述

航拍圖像和遙感圖像識別長期以來得到了測繪、國防、農業等領域的廣泛應用，隨著深度學習和卷積神經網絡（Convolutional Neural Networks，CNN）的高速發展，地表覆蓋物分類等基于遙感影像的自動化分析技術也得到了一定程度的發展。

在過去的10年內，視覺模型的發展主要有兩個主流方向，分別是增加模型對于圖像的理解能力和降低模型理解圖像的算力消耗。通過增加模型復雜度提高精度已經得到了充分研究，但如何在不損失太大精度的前提下，降低模型的大小和算力消耗仍有待探索。為了使深度學習模型得到更廣泛的應用，尤其是在移動端設備上應用，這樣的探索是必須的。

目前，學術界已有一些模型優化的成果，比如MobileNetV2[1]和ShuffleNetV2[2]。從優化的不同時間階段，模型優化可以分為訓練階段和推理階段。在模型訓練階段，可以更改模型權重更新的過程，從而提高訓練效率;在模型推理階段，可以將訓練好的模型進行編譯優化，得到一個體積更小、運行更快的部署模型。

（1）訓練階段的優化

l空間可分離卷積：在不考慮通道維度的情況下，通過將卷積核分解降低運算量，但并不是任何卷積核都可以分解為兩個向量，因此不具備通用性。而且，計算矩陣分解還需要消耗一定的算力。

l深度可分離卷積：在考慮通道維度的情況下，將矩陣連乘的過程優化為乘法的加法。需要注意的是，深度可分離卷積是減少卷積的參數量的近似算法，和樸素的卷積并不是嚴格等價的。

l量化：在不過分影響模型效果的前提下，將部分32位浮點數精度的數值降低為16位浮點數和8位整數，這樣可以減少模型體積、加快訓練的速度。

l剪枝：在不過分影響模型效果的前提下，將模型的部分結構刪除，從而減少模型體積、加快推理速度。

l聚類：通過共享權重的方式來減少權重的數量，這樣可以減少模型體積、加快模型收斂速度。在樣本稀疏或者樣本分布不均衡的部分場景下，可以提升模型的效果。聚類對于模型效果和收斂速度的提升都是鑒于當前的深度學習模型的參數數量很大、擬合能力很強這一事實。

（2）推理階段的優化

l常量折疊：在編譯階段將一些常量表達式的值計算出來，然后寫入二進制文件中（在模型編譯優化的場景下，就是優化后的模型），同時也是許多現代編譯器使用的編譯優化技巧。這樣可以通過編譯時的單次計算換取運算時的若干次重復運算。并且，由于把常量表達式替換為常量，也減少了不少空間的存儲，因此也可以減少模型的體積。

l減少死區代碼：經過模型的優化之后，可能有許多的計算步驟被跳過了，因此會出現大量的不會被觸達的指令區域。減少死區代碼的方法是，通過遍歷所有指令來將這些代碼移除掉，達到減少模型體積的效果。

l算子融合：由于現代的機器學習建模框架都是以微型算子與圖的形式來表示一個模型。適當地將一些常見的算子拓撲的序列融合成一個算子可以極大地減少算子的內存占用和算子之間的數據傳輸壓力，這樣就會達到減少模型大小和加快模型計算速度的目的。

訓練階段廣泛被使用的框架有TensorFlow[3]和PyTorch[4]等，推理階段有ARM NN和TensorRT等。雖然這些框架都具有一定的跨平臺特性，但要想做到對所有的目標設備（CPU、GPU、NPU、TPU）都能實現很好的推理性能是很難的。這是因為，各種硬件設備的計算特性有所差異，所以框架之間并不存在強通用性。

本文介紹了一些重要的深度學習模型優化方法，并基于遙感影像語義分割問題進行了實驗。實驗結果顯示，本文介紹的方法在不損失太多精度的前提下，對于縮小模型大小、提高訓練和推理效率有著很好的效果，有利于遙感影像語義分割算法在移動端的部署與應用。

2 方法

本文介紹神經網絡訓練和推理階段的若干優化方法，具體的流程如圖1所示。在訓練階段，引入混合精度、權重剪枝、權重聚類獲得模型1;在推理階段，通過編譯優化獲得模型2，即最終模型。

2.1 混合精度

如今，大多數深度學習模型使用的數據精度是32位浮點數，占用4字節內存。但是，使用16位浮點數可以達到更快的運算速度和更小的內存占用，因此產生了混合精度技術。混合精度是訓練時在模型中同時使用16位和32位浮點類型，從而加快運行速度、減少內存使用的一種訓練方法。讓模型的某些部分保持使用32位類型的目的是，既能保持數值穩定性，又能縮短模型的單步用時，在提高訓練效率的同時仍可以獲得同等的模型精度。

在混合精度訓練時，由于權重、激活值、梯度等都使用16位存儲，會更容易產生舍入誤差和數據溢出的問題。為了使模型的準確性不損失過多，可以使用權重備份和損失縮放的方法[5]。權重備份是指，在優化步驟中，維護、更新一份32位的權重。在每次迭代中，32位權重被轉為16位，用于前向傳播和反向傳播。權重備份主要防范的問題是，當產生過小的梯度，乘學習率后很容易下溢為0，對模型準確性產生不利的影響。使用權重備份后，可以通過備份的32位權重恢復準確性。損失縮放的方法啟發于網絡訓練期間對數值范圍的觀察，大部分16位可表示范圍未被使用，而許多低于可表示最小值的值變成了0。如果將這些數值增大到2的冪次倍，可以避免下溢產生的誤差。最有效的方法是，將前向傳播后得到的損失值縮放，再進行反向傳播，根據鏈式法則，后續的梯度也會被縮放同樣的倍數。在選取縮放倍數的時候，只要不導致溢出，縮放因子越大越好。如果產生溢出，將使權重或梯度產生無窮大或NaN，而這是無法恢復的。

2.2 權重剪枝

權重剪枝通過刪除神經網絡層之間的連接，從而減少參數，進而減少計算中涉及的參數和運算量，有利于模型壓縮，對于模型的存儲與傳輸來說都有好處。權重剪枝方法可以消除權重張量中作用不大的值，這一步通常會用到某些啟發式方法，比如計算權重張量的平均激活值來對權重張量中的值進行優先級排序，找到一定范圍內的低優先級的權重，再通過將神經網絡中這些參數的值設置為0，以消除神經網絡各層之間不必要的連接。由于在權重剪枝過程中，如果一次性剪枝掉太多張量中的值，可能會對神經網絡造成不可逆的損傷，因此，權重剪枝方法常常是迭代式進行的，即訓練與剪枝過程的交替重復。稀疏張量可以實現很好的壓縮效果，因此很適合進行壓縮。權重剪枝過程之后，由于張量中存在大量為0的參數，所以可以很好地進行模型壓縮，以達到降低模型大小的目的。

下面敘述一個具體的剪枝方法。對于選擇修剪的每一層，添加一個二進制掩碼變量，與每層的權重形狀相同，用于確定哪些權重參與運算。前向傳播時，對權重按絕對值進行排序，絕對值最小的權重對應的掩碼被設為0，直到達到所需的稀疏度。反向傳播過程中，在前向傳播中被屏蔽的權重不會得到更新。需要注意的是，為了避免對模型的準確性產生太大影響，設定的稀疏度需要經過多次調整。學習率大小與其調整方式是另一個對模型精度影響很大的地方。如果學習率會下調很多，在學習率大幅下降、網絡已經被剪枝后，網絡將難以從剪枝帶來的準確率下降中恢復。如果初始學習率過高，可能意味著當權重尚未收斂到一個好的解時修剪權重。這里描述的剪枝方法不依賴于任何特定網絡的屬性，具有高度的泛化性。

2.3 權重聚類

權重聚類通過用相同的值替換權重中的相似權重來減小模型的大小。通過在模型的訓練權重上運行聚類算法，可以找到這些相似的權重，并把這些相似的權重用同樣的值進行代替。權重聚類在減少模型存儲空間和傳輸大小方面具有直接優勢，因為具有共享參數模型的壓縮率要比不具有共享參數的模型高得多。這樣的模型壓縮方法類似于權重剪枝，不同之處在于，權重聚類是通過增加共享權重的參數個數來實現的，而剪枝是通過將低于某個閾值的權重設置為0來實現的。由于權重聚類會將模型中所有參數替換為給定個數的共享權重值，所以在應用權重聚類方法之前，通常需要對模型進行一些預訓練，使得模型的參數更新到一個可接受的水平之后，再使用權重聚類對參數進行微調。將模型進行權重聚類后，可以應用常用的壓縮算法進行壓縮。由于權重聚類后的模型中存在大量的共享參數，可以實現更高的壓縮率。

在具體實現過程中，由于聚類中心對于聚類效果有影響，需要選擇合適的權重初始化方法，根據文獻[6]，線性初始化是最優的方法。計算相似權重的常用方法是K-means，并用歐氏距離度量權重和聚類中心的距離。聚類完成后的聚類中心就是共享權重。

2.4 編譯優化

編譯優化是將編譯前端生成的模型在編譯器中端通過計算圖的優化方法對模型進行優化。需要注意的是，針對不同的目標設備，使用的優化方法以及優化策略是有所不同的，以確保生成優化后的計算圖能夠在目標設備上表現出最佳性能。

編譯優化的方法主要有：運算符合并，將多個小運算（加、和）合并為一個運算;常量折疊，可以預計算能夠確定的部分計算圖;靜態內存規劃，可以預分配臨時變量占用的空間，避免多次內存申請和釋放;數據分布變換，可以將數據分布變換為有利于進行高效后端運算的分布，比如將行主序或列主序的矩陣變換為更適合目標硬件的存儲形式。

3 實驗與分析

3.1 實現細節

本文以遙感影像地表覆蓋物分類為例對所介紹的方法進行驗證。遙感影像數據的來源為上海市某地區的航拍圖，并按照建筑、水系、植逐像素地進行了人工標注。完成語義分割任務的基準深度學習網絡是深度為4的U-Net[7]，并融合了MobileNetV2[1]，使用TensorFlow實現。輸入網絡的圖片大小均為1024*1024。訓練批次大小設為4，迭代20代，損失函數使用Lovasz Loss[8]。

混合精度使用TensorFlow的Keras API實現，權重剪枝和權重聚類使用TensorFlow提供的工具包Pruning API和Weight Clustering API實現，編譯優化使用TVM[9]實現。在模型進行儲存時，權重仍然按照與權重剪枝或權重聚類前的模型一樣的Tensor格式進行存儲，而并非會因為權重矩陣變得更為稀疏而采用稀疏矩陣的方式進行存儲。所以，為了減少模型權重矩陣中的冗余數據的存儲，在存儲模型時采用gzip算法進行壓縮。

3.2 實驗結果

混合精度實驗的結果如表1，有預訓練時混合精度會縮短訓練時長并降低Loss，無預訓練時則相反。使用混合精度后，縮短訓練時長的表現并不穩定，并且會在一定程度上影響訓練結果，但內存占用被大幅減少了。

權重剪枝與權重聚類的實驗結果如表2，其中耗時表示一張圖片的推理耗時。可以看出，權重剪枝與權重聚類對于推理耗時的影響非常有限，但是能大大縮減模型大小，而且對輸出準確度沒有負面影響。之所以對推理耗時幾乎沒有影響，是因為盡管這兩種策略會提高矩陣的稀疏度，但是由于TensorFlow默認存儲的方式仍然是Dense Tensor，所以在進行計算時，沒有對稀疏張量提供計算加速的支持，仍然采用默認的計算方法。

權重剪枝與權重聚類后對一張圖的預測結果如圖2所示，從左向右分別為：原圖、原預測圖、應用權重剪枝的預測圖、應用權重聚類的預測圖。可以看出，權重剪枝與權重聚類對于預測結果沒有太大的影響，但權重剪枝會造成比權重聚類更多的誤分類。

使用TVM進行編譯優化實驗的結果如表3，對Conv2D算子進行圖優化。可以看出，經過多輪優化，推理時間逐步減少。

4 結論與展望

本文針對遙感圖像語義分割在移動平臺應用的問題，研究了神經網絡優化算法。本文探索了混合精度、權重剪枝、權重聚類、編譯優化的方法，在開源平臺上進行了實驗，分析了各種方法的優勢與不足。混合精度能夠大幅降低內存使用，但對于訓練時間和訓練效果的影響不可預料;權重剪枝和權重聚類能大幅度減小存儲模型的大小，對推理時間的影響幾乎可以忽略不計;編譯優化和平臺具有很高的相關性，難以部署。為了進一步探索優化算法的適用性，可以在不同的網絡模型上進行實驗，在不同平臺上測試性能表現。

參考文獻：

[1] Sandler M，Howard A，Zhu M L，et al.MobileNetV2：inverted residuals and linear bottlenecks[C]//2018 IEEE/CVF Conference on Computer Vision and Pattern Recognition.Salt Lake City，UT.IEEE，2018： 4510-4520.

[2] Ma N N，Zhang X Y，Zheng H T，et al.ShuffleNet V2：practical guidelines for efficient CNN architecture design[C]//Computer Vision – ECCV 2018，2018： 116-131.

[3] Abadi M， Agarwal A， Barham P， et al. Tensorflow： Large-scale machine learning on heterogeneous distributed systems[J]. arXiv preprint arXiv：1603.04467， 2016.

[4] Paszke A， Gross S， Massa F， et al. Pytorch： An imperative style， high-performance deep learning library[J]. Advances in neural information processing systems， 2019（32）： 8026-8037.

[5] Micikevicius P， Narang S， Alben J， et al. Mixed precision training[J]. arXiv preprint arXiv：1710.03740， 2017.

[6] Han S， Mao H， Dally W J. Deep compression： Compressing deep neural networks with pruning， trained quantization and huffman coding[J]. arXiv preprint arXiv：1510.00149， 2015.

[7] Ronneberger O，Fischer P，Brox T.U-net：convolutional networks for biomedical image segmentation[C]//Medical Image Computing and Computer-Assisted Intervention-MICCAI 2015，2015：234-241.

[8] Berman M，Triki A R，Blaschko M B.The lovasz-softmax loss：a tractable surrogate for the optimization of the intersection-over-union measure in neural networks[C]//2018 IEEE/CVF Conference on Computer Vision and Pattern Recognition.Salt Lake City，UT，USA.IEEE，2018：4413-4421.

[9] Chen T，Moreau T，Jiang Z，et al.TVM：An automated end-to-end optimizing compiler for deep learning[C]//Proceedings of the 13th USENIX Symposium on Operating Systems Design and Implementation（OSDI 8）.Carlsbad，USA，2018：578-594.

收稿日期：2022-02-10

作者簡介：王曉鋒（1982—），男，上海人，高級工程師，碩士，主要研究方向為土木工程;張猷（1982—），男，上海人，高級工程師，本科，主要研究方向為輸變電工程建設管理;夏嚴峰（1990—）男，上海人，工程師，碩士，主要研究方向為輸變電項目管理;王森（1999—），男，黑龍江慶安人，本科，主研方向為圖像分割;王輝（1997—），男，安徽廬江人，碩士在讀，主要研究方向為圖像分割。