基于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%，为在嵌入式环境中实现遥感影像场景分类以及后续的智能解译提供了一种有效的思路。