基于ResMLP模型的遙感影像場景分類算法式實現

周浩楠,夏魯瑞,李 森

(1.航天工程大學 研究生院，北京 101416；2.航天工程大學 某研究中心，北京 101416)

1 引言

隨著衛星對地觀測技術的不斷發展，遙感影像數據的獲取不斷向多源化、全球化方向發展，同時這對遙感影像智能解譯提出了更高的要求。遙感影像場景分類是遙感影像智能解譯的重要環節，它經常被用作衡量對遙感影像解譯程度的基準，其進展易于進一步推動遙感影像目標檢測或分割任務的改進。

近年來的大多數遙感影像場景分類研究都是使用深度卷積神經網絡(CNN)來實現，主流CNN模型有GoogleNet[1],ResNet[2],DenseNet[3]。但是圖像分類中使用的CNN方法受益于多年發展而來的調參和優化方法[4]。與主流CNN模型相比，基于transformer的模型不使用這些復雜的技巧，就能在視覺任務上表現出良好的性能。transformer[5]是一種主要基于自注意力機制構建的網絡結構。transformer廣泛應用于自然語言處理領域，例如著名的BERT[6]和GPT3[7]模型。受transformer在自然語言處理方面發展的啟發，研究人員最近將transformer應用于計算機視覺任務，如圖像識別、對象檢測和圖像處理。DETR[8]將目標檢測視為一種集合預測問題，并使用transformer構建編碼器-解碼器架構來解決。IPT[9]利用transformer在一個模型中處理多個視覺任務。這些基于transformer構建的模型在各種計算機視覺任務上表現出了良好的性能。

與CNN方法相比，vision transformer(VIT)[10]沒有使用卷積結構固有的假設和平移等變性。VIT在不使用大量優化和調整下，實現了與CNN最先進模型相近的性能，是當前倍受矚目的模型。但是當圖像分辨率較大時(例如遙感影像)，VIT對圖像的特征提取的速度會變慢。因為VIT中的自注意力機制會占用大量內存，并且對圖像尺寸的變化不敏感。VIT訓練階段需要大量算力和大量時間來進行有效地學習，并且這類模型通常十分龐大且內部結構復雜，嵌入式系統難以提供硬件層面的加速支持，直接使用VIT作為主干網絡的算法架構在嵌入式系統中實現在軌智能場景分類十分困難。

最近，來自Facebook AI的學者們提出了一種主要由MLP構成的算法架構ResMLP[11]：類似于VIT直接將圖片切分為不重疊的patch作為輸入，但是不使用任何形式的注意力機制，僅僅包含線性層與 GELU[12]非線性激活函數，整體結構類似于VIT但更加簡單。這種簡單的基于MLP構造的簡單模型在ImageNet數據集中[13]取得了與當前性能最佳的CNN模型、VIT模型相當的性能。同時，基于MLP的線性算法架構具有更強的解釋性，易于獲得嵌入式系統的加速支持，比VIT更適合搭載在嵌入式環境中運行。本文中選用ResMLP進行遙感影像場景分類，并且成功將模型移植到了Jetson TX2中，實現了85.9%的TOP-1準確率。

2 基于ResMLP模型的遙感影像場景分類算法

2.1 算法流程

ResMLP首先將輸入的遙感影像劃分為n×n個不重疊的圖像分塊patch，其中n可以按照需求設定。然后將這n×n個不重疊的遙感影像分塊輸入進一種線性層“殘差MLP層”中，得到一組維度為d，共n×n個嵌入向量。隨后將這組d維嵌入向量再次輸入到一組殘差MLP層，并將殘差MLP層的輸出嵌入整合為一維向量。這個一維向量中含有表征遙感影像的特征信息。最后將該一維向量輸入到分類器中以完成對輸入遙感影像的場景分類。整體處理流程示意圖如圖1所示。

圖1 基于ResMLP的遙感影像場景分類算法流程

2.2 ResMLP模型

ResMLP模型由殘差MLP層和ResMLP分類器組成。輸入遙感影像patch經過殘差MLP層提取遙感影像中的特征信息，將特征信息輸入ResMLP分類器中進行場景分類。

2.2.1殘差MLP層

ResMLP包含數個殘差MLP層。殘差MLP層的結構為一個線性層后接一個前饋層，類似深度殘差網絡ResNet，因此命名為殘差MLP層。殘差MLP層使用如下仿射變換對由輸入的遙感影像切割得到的patch進行縮放以及位置變換，實現對遙感影像特征的提取：

Affα，β(x)=Diag(α)x+β

(1)

Z=X+Aff((PAff(X)T)T

(2)

Y=Z+Aff(QGELU(MAff(Z)))

(3)

其中α、β、P、Q、M為可學習參數。上述3個公式在推理過程中屬于線性層的一部分，因此殘差MLP層在推理過程中不會占用額外的算力資源。使用獨立于數據的殘差MLP層替換具有數據依賴性的自注意力機制，降低模型的訓練難度。不同于卷積神經網絡的對局部特征進行提取并且進行權值共享，殘差MLP層對各個patch進行特征提取時不進行patch之間的權值共享，并且不像VIT一樣使用位置嵌入來確定patch的位置，減少了整體結構的運算量和推理時間。

2.2.2ResMLP分類器

不同于VIT使用tokens構建復雜的分類器，本文中的ResMLP中使用結構簡單的線性分類器，將遙感影像patch生成的嵌入表示和類別標簽連接。因為殘差MLP層在運算過程中對獨立的每個patch分別進行特征提取，因此本文中算法在線性分類器前加上了對鄰域信息進行偏重考慮的average pooling來兼顧局部和整體特征信息，使整合提取到的特征更為準確。基于這些設計，ResMLP的整體結構相比VIT更為精簡，ResMLP對輸入的遙感影像的分類速度比VIT更快。

3 算法嵌入式實現

將基于ResMLP的遙感影像分類算法遷移到嵌入式系統中，首先需要在服務器端訓練基于ResMLP的遙感影像場景分類模型。模型在嵌入式系統中充分發揮性能需要合適的算法運行環境，因此本文中對嵌入式系統的算法環境進行了配置。隨后將算法和模型遷移到嵌入式系統NVIDIA Jetson TX2中，最后對遷移得到的算法的性能進行評估。遷移算法的流程示意圖如圖2所示。

圖2 基于ResMLP的遙感影像場景分類模型嵌入式移植流程

3.1 嵌入式系統軟硬件環境搭建

NVIDIA Jetson TX2能夠在低功耗下為邊緣端AI計算提供出色的性能。最低功耗僅有7.5 W，最高時也僅有15 W，完全可以滿足軌道處理衛星的能耗限制。其CPU為NVIDIA Denver 2 64位CPU與四核Arm?Cortex?-A57 MPCore復合處理器，內存為8 GB 128-bit LPDDR4，GPU為NVIDIA Pascal架構，配有256個NVIDIA CUDA核心以及59.7 GB/s帶寬的顯存以及各種標準硬件接口。

在進行移植前，使用JetPack 4.4對Jetson TX2進行刷機，使用的操作系統為Ubuntu 18.04.5 LTS(GNU/Linux 4.9.140-tegra aarch64)，CUDA工具包版本為10.2。使用的Python的版本為3.6.9，pytorch的版本為1.7.0，torchvision的版本是0.8.0。

3.2 模型訓練

本文中使用的數據集為MASATI dataset[14]。該數據集由阿利坎特大學于2018年提出，其中包含動態海洋環境中的彩色圖像。數據集含有7類共7 389張遙感影像，這些遙感影像均從必應地圖中獲取，平均尺寸為 512×512。MASATI數據集的部分圖片如圖3所示。

圖3 MASATI數據集部分圖片

ResMLP遙感影像分類模型訓練環境為：Windows 10操作系統，CPU為Intel i9-10900k,GPU為NVIDIA-3090 24GB，內存為64GB。訓練過程中batch size取512，優化器為AdamW[15-16]，學習率取值0.002 5，weight decay取0.01，網絡的輸入尺寸設定為224，num workers取4，總共訓練300個epochs。訓練過程中學習率以及準確率的部分記錄如圖4和圖5所示。

圖4 ResMLP遙感影像分類模型1～60 epoch訓練記錄

圖5 ResMLP遙感影像分類模型170～220 epoch訓練記錄

從訓練記錄可以看出，ResMLP遙感影像場景分類模型的收斂速度很快，在第30輪訓練中就得到了接近最高準確率的性能，在第61輪訓練時得到了最高性能。同時也可以注意到，模型訓練后期會出現性能下降的問題。由此可知，ResMLP遙感影像場景分類模型訓練的輪數不需要超過100輪就可以得到最佳模型，訓練速度快于需要在大規模數據中進行預訓練的VIT。

3.3 嵌入式移植

3.3.1算法環境配置

模型訓練完成后，將訓練得到的ResMLP遙感影像場景分類模型移植至TX2中，并且根據Windows 10下的訓練和測試環境進行TX2的環境配置。由于TX2基于arm架構，圖像分類算法中的部分依賴庫是基于x86架構的，無法直接在TX2中進行安裝，例如h5py，lmdb，scikit-image等。本文中選擇通過交叉編譯的方式，將x86架構下的依賴庫編譯為適應arm架構的依賴庫以實現環境配置。

以scikit-image的解決方案為例，安裝scikit-image需要先安裝其運行的關聯庫，其中要求matplotlib的版本不低于1.3.1，numpy版本不低于1.7.2，scipy版本不低于0.9.0，six版本不低于1.7.3，networkx版本不低于1.8，pillow版本不低于2.1.0，dask版本不低于0.5.0。其中networkx，dask，scipy這三個庫在使用pip安裝時出現錯誤，表明環境中仍舊缺少依賴庫。經過閱讀這3個依賴庫的源碼得知，這3個庫是在x86架構下運行的，無法直接在基于arm架構的TX2上直接安裝運行，此處通過交叉編譯解決：

sudo apt-get install libblas-devliblapack-dev

對依賴庫進行交叉編譯之后再次安裝scikit-image，此時scikit-image安裝成功。其他遇到這種安裝問題的依賴包，本文中均按照類似的方法進行安裝配置。

3.3.2算法代碼修改

在Ubuntu 18.04中完成環境配置后，需要根據實際環境對算法中的部分代碼進行調整，以適應嵌入式系統中的環境，完成對算法的移植。

首先需要調整算法的計算精度。NVIDIA turning架構對Float16和Int8的矩陣乘法有額外加速。雖然Float16存在計算溢出和舍入誤差，但是相比Float32其內存占用少和計算更快，能夠大幅加快算法推理速度。本文中在訓練模型時使用了Float16和 Float32混合精度的訓練策略，即在內存中用Float16做儲存和乘法使得計算獲得加速，用Float32做累加避免舍入誤差。采用混合精度訓練可以使訓練時間大大減少，它的缺陷是只能在支持Float16操作的一些特定類型的顯卡上面使用，而且會存在溢出誤差和舍入誤差。在將本文中的算法移植至NVIDIA Jetson TX2時，考慮到TX2對Float16有很好的加速，因此將推理代碼中所有使用Float32精度的部分全部改為使用Float16精度進行計算。

數據的讀取方式也需要進行修改。訓練模型時的環境為Windows 10，而NVIDIA Jetson TX2的運行環境為Ubuntu 18.04系統，二者在對數據進行讀取時的算法編寫格式存在差別。Windows 10中的數據路徑以“/”隔開，并且會使用“:”，這些在Ubuntu 18.04系統直接進行運行會報錯，因此算法在移植過程中需要對數據的讀取語句進行修改：將“/”用“”代替，同時去掉所有路徑中的“:”，并將移植前的文件存儲路徑修改為新的存儲路徑。

3.4 運行結果

經過上述算法遷移過程，基于ResMLP的遙感影像場景分類模型成功移植到了NVIDIA Jetson TX2中并成功運行，完成了算法的嵌入式實現，如圖6所示。

圖6 ResMLP遙感影像場景分類算法嵌入式移植結果

由圖6可知，基于ResMLP的遙感影像場景分類模型的top-1準確率為85.9%，param值為18.4 M，模型的大小為210 MB。同時本文中也在相同環境下對VIT進行了移植。將2種算法從模型推理所需時間以及top-1準確率2個指標進行了對比，結果如表1所示。

表1 VIT與ResMLP算法性能對比

由表1可知，移植在TX2中的基于ResMLP的遙感影像場景分類模型對單張512×512遙感影像進行場景分類平均需0.1 s，可以達到85.9%的準確率。在TX2環境中該模型的準確率和處理速度均優于VIT。

4 結論

本文中提出了一種基于ResMLP的遙感影像場景分類算法，并且在嵌入式系統NVIDIA Jetson TX2中進行了嵌入式實現。由移植結果可以看出，本文中提出的基于ResMLP的遙感影像場景分類算法能夠在嵌入式系統NVIDIA Jetson TX2實現并且準確快速的對遙感影像進行場景分類。此算法訓練模型所需時間短，并且在運行過程中對計算機資源的占用小，模型的推理速度快于VIT，對遙感影像場景分類的top-1準確率能夠達到85.9%，為在嵌入式環境中實現遙感影像場景分類以及后續的智能解譯提供了一種有效的思路。