基于对比自监督的遥感图像场景分类

吴奕恒,陈知明,戴颖成,欧阳文欣,李建军,刘 峰

(中南林业科技大学 计算机与信息工程学院,湖南 长沙 410000)

0 引言

遥感图像场景分类是遥感图像信息处理的关键组成部分,旨在将场景图像划分为不同的语义类别,目前已被广泛应用于土地资源管理[1]、城市规划[2]等领域。

基于深度卷积神经网络(Deep Convolutional Neural Networks, DCNN)的深度学习方法[3-5],因其强大的层次特征提取能力已成为遥感图像场景分类领域的研究热点之一。然而,训练DCNN 需要大量的标注图像数据,且对数据标注是一项繁琐的工作并且代价昂贵[6]。Gidaris 等人[7]提出基于Rotation 的ConvNet 模型,能自动生成旋转图像和旋转角度标签,从而实现使用无标注数据构造辅助任务进行自监督学习。但旋转角度标签仅能提供间接、可量化的信息,无法有效表达类内差异性等复杂信息。针对该问题,Bao 等人[8]提出了基于掩码图像建模的BEIT 模型,将图像分割成块后生成高维嵌入序列,随机遮盖序列的部分并构造嵌入序列重建任务进行自监督学习。然而,生成嵌入序列的过程丢失了大量特征信息,无法有效利用遥感图像丰富的潜在语义。为此,He 等人[9]提出掩码自编码器(Masked Autoencoders,MAE)模型,将无标记图像分割为等大小图像块,随机遮盖后逐像素重建原始信号以提取图像的潜在特征表示。MAE 模型在进行特征提取的过程中,进行了图像分割与随机遮盖,导致上下文信息与图像局部特征的丢失,当面对地物建筑相似、整体布局不同的对象时出现误分的现象,从而导致分类精度的降低。

为了解决MAE 模型在特征提取过程中局部语义上下文信息丢失的问题,本文在MAE 的基础上进行改进,提出一种掩码重建下的对比学习自监督遥感图像场景分类方法——对比掩码自编码器(Masked Autoencoders with Contrast,MAECO),一方面,通过掩码重建任务强调遥感图像当中的全局特征;另一方面,结合对比学习同时构建两个负样本队列,通过对比正负样本而强化局部特征,补充掩码重建中丢失的局部语义信息。

1 MAECO 模型

MAE 在进行编码解码的过程中,存在着局部上下文语义信息丢失的问题,因此,本文提出了一种自监督学习框架MAECO,其网络结构如图1 所示。MAECO 是一种3 阶段的自监督模型,分两个自监督预训练阶段和一个微调阶段,分别为掩码重建预训练阶段,对比学习预训练阶段和微调阶段。该模型在进行掩码重建预训练的过程中,侧重加强提取全局语义信息的能力。在对比学习预训练过程中,MAECO 侧重强化对局部特征的捕捉能力。在微调阶段,MAECO 与具体下游任务进行适配。

图1 MAECO 总体结构Fig.1 The structure of MAECO

MAECO 具体流程如下:第一阶段,使用无标签的ImageNet 数据集对MAE 进行掩码重建任务预训练;第二阶段,提取MAE 的编码器,构造MAECO 编码器和两个冲量MAE 编码器,在无标签的遥感图像数据集上进行对比学习;第三阶段,提取训练好的MAECO 编码器,增加线性分类器和SoftMax 层以用于特征分类,在有监督的遥感图像数据集上进行微调。下面详细介绍每阶段工作。

1.1 掩码重建预训练

MAE 模型由MAE 编码器和掩码重建解码器构成,如图1 掩码重建预训练部分所示。MAE 编码器为24 层视觉Transformer,掩码重建解码器为8 层视觉Transformer,结构如图2 所示。

图2 MAE 结构Fig.2 The structure of MAE

在将无标签原始图像输入MAE 模型前,首先按75%的比例对原始图像进行遮盖。然后,将图像分割为图像块,仅将未被遮盖的图像块送入MAE 编码器生成对应的特征图,以使用较少的计算资源训练较深的视觉Transformer 模型,提高模型的特征提取能力。随后,将这些特征图和被遮盖的图像块按在原始图像中的位置顺序排列,送入掩码重建解码器,进行逐像素的重建任务。在进行掩码重建任务时,常使用SSIM 损失函数及其变体[10]。然而,这类损失函数计算复杂度过高,会占用大量计算资源,使训练时间大幅延长。故使用简单的MSE 损失函数,公式如下:

其中,W为图像宽度,H为图像高度,为(i,j)位置的像素真实RGB 值,为(i,j)位置的像素预测RGB值。使用经掩码重建预训练的MAE 模型进行图像重建的效果如图3 所示。由图3 可知,掩码重建训练后,该模型能够显式提取图像的色彩特征和主要纹理特征等全局特征,但忽视了局部特征(如机场中停放的飞机在重建时丢失)。

图3 图像重建效果Fig.3 Image reconstruction results

1.2 对比学习预训练

掩码重建预训练将图像随机遮盖后提取全局特征信息,造成了局部特征信息的丢失。而遥感图像需要局部语义位置信息等局部特征才能进行精准分类。为此,本文提出对比学习模块,进行对比学习预训练,以强化模型的局部特征提取能力,其中对比学习结构如图4 所示。

图4 对比学习模块结构Fig.4 The structure of contrastive learning module

为了学习更多的潜在的局部一致性特征信息,同时使用两组不同的负样本分别进行对比学习。由于丰富的数据增强模式可以迫使模型学习潜在的不变特征,分别对图像进行两组不同类型数据增强,一组进行色彩组合增强,即随机裁剪、随机翻转、色彩抖动、随机高斯模糊,另一组进行灰度组合增强,即随机裁剪、随机翻转、灰度化、随机高斯模糊,生成的增强图像分别输入色彩抖动冲量MAE 编码器和灰度化冲量MAE 编码器,生成对应的特征图。使用相似的数据增强模式,可以在丰富图像表征多样性的同时,保留图像的潜在一致性,使网络可以学到潜在的共同特征。

为了减少对内存资源的占用,本文使用两个队列存储数据增强特征图,利用队列“先进先出”的特性,将原图和队尾的数据增强特征图(即新生成的数据增强特征图)作为正样本。由于使用两个队列存储负样本,若使用反向传播算法更新两个冲量MAE 编码器参数,计算复杂度较高,会占用大量计算资源,因此需要一种新的方法来更新参数。结合冲量更新法[11]和Adam 算法更新参数的思路,提出一种新的参数更新方法:

其中,θMAECO,θJitter,θGrey分别为MAECO 编码器、色彩抖动冲量MAE 编码器和灰度化冲量MAE 编码器的模型参数。β为冲量更新参数,决定了模型参数更新的快慢,默认为0.999。MAECO 编码器使用正常的反向传播算法进行参数更新,而两个冲量MAE 编码器利用其它编码器协助进行参数更新。

本文使用InfoNCE[12]作为单个对比学习的损失函数:

其中,q为原图对应特征图,k+(即kK)为正样本特征图,ki为队列中第i个样本,K为队列中不包括正样本在内的特征图个数。τ为温度参数,决定了对比损失对困难负样本的关注程度,越大则对困难负样本越关注,默认为0.07。

两组对比学习的复合对比损失函数为

其中,LJitter和LGrey分别为色彩抖动对比学习和灰度对比学习的InfoNCE 损失函数,α为权重参数,决定了增强图像组的对比学习对模型更新的影响权重,默认为0.5,表明两组对比学习对模型更新的影响权重相当。

1.3 微调

在微调阶段,MAECO 编码器迁移为下游任务网络的主干,添加线性分类器和SoftMax 层以输出分类结果,结构如图1 微调部分所示。

在有监督遥感图像数据集上进行微调时,使用交叉熵损失函数:

其中,N为样本数,M为类别数,yij为第i个样本是否为第j类(是则取1,否则取0),pij为第i个样本预测为第j类的可能性。

2 实验结果与分析

2.1 数据集

本节使用AID 数据集[13]和NWPU-RESISC45数据集[14]这两种公开遥感图像数据集进行实验。以上数据集的具体信息如表1 所示,样本示例如图5 所示。

表1 数据集详情Tab.1 Table of dataset details

图5 AID 数据集和NWPU-RESISC45 数据集遥感场景示例Fig.5 Samples of AID dataset and NWPU-RESISC45 dataset

为与不同文献的方法实验进行公平比较,根据文献[13-14]中的Benchmark 设定进行训练集和测试集比例设置,即对AID 数据集的每种类别按比例进行随机抽样,取20%的图像作为训练集,其余80%的图像作测试集。此外,增设训练比例为50%的AID 数据集进行额外对比实验。对于NWPU-RESISC45,按10%和20%的比例随机抽取作训练集和测试集。

2.2 评估标准

本文使用总体分类精度(Overall Accuracy,OA)和混淆矩阵(Confusion Matrix)作为评估模型性能的指标。

总体分类精度计算分类正确的样本占总样本的比例,反应了整体分类性能。总体分类精度计算公式如下:

其中,N为总样本数,M为类别数,ri为第i类预测正确的样本数。

混淆矩阵每一行代表样本的真实类别,每一列代表预测类别,整个矩阵反映了模型对每个样本类别的预测水平(由于生成混淆矩阵时会有取整进位操作,故混淆矩阵部分行之和不为1)。

2.3 实验环境和参数设定

实验软件环境为Ubuntu 16.04,Python 3.6,深度学习框架为PyTorch 1.8。硬件环境为Intel i9-79200X CPU,NVIDIA RTX 2080 Ti 显卡,32G内存,11G 显存。

实验训练使用随机梯度下降优化器,批量大小为32,所有阶段迭代次数为100,初始学习率设为0.001,权重衰减系数设为0.05。为确保准确性,每组实验均独立进行5 次。

2.4 对比实验

在AID 数据集和NWPU-RESISC45 数据集上将MAECO 方法和当前主流自监督学习方法进行对比,实验结果如表2 和表3 所示。同时,增设与主流监督方法对比实验,实验结果如表4 所示。

表2 不同自监督方法的精度比较Tab.2 Accuracy comparison of different self-supervised methods%

表3 不同方法训练时间和模型参数比较Tab.3 Comparison of model parameters and training time of different methods

表4 MAECO 和全监督方法精度比较Tab.4 Accuracy comparison of MAECO and supervised methods%

由表2 可知,MAECO 的总体精度全方面超越了现存自监督学习方法,和次优模型MAE 相比,在两种训练比例的AID 数据集上精度提高6.22%和2.47%,在NWPU-RESISC45 数据集上精度提升9.24%和4.24%。这是由于MAECO 模型使用两段式自监督预训练将自注意力机制和对比学习结合,能够最大程度地提取遥感图像的全局特征和局部特征,对图像的语义信息利用率更高。从表4可以看出,MAECO 在两种比例设置的AID 数据集上相较最优方法MG-CAP(Sqrt-E)总体精度下降2.43%和0.75%,优于文献[14]、[18] 的方法,接近主流全监督遥感图像场景分类方法,在NWPURESISC45 上的预测精度已超过当前主流全监督方法。由表3 可知,MAECO 训练时长在MAE 和MoCo之间,参数量与MAE 相当,在精度和运行时间之间取得了平衡。由于使用视觉Transformer,MAECO相比全卷积神经网络参数量更多,但通过权值共享的方法有效压缩模型参数,使其与MAE 相当。

为更好地展现MAECO 的性能,绘制了对应混淆矩阵。图6～7 展示了MAECO 在AID 数据集和NWPU-RESISC45 数据集上的混淆矩阵。在AID数据集中机场(1:Airport)、山脉(16:Mountain)、港口(21:Port)和高架桥(30:Viaduct)等纹理差异性小、包含丰富显著性特征的场景分类精度较高,中型住宅区(15: MediuResidential)和度假村(23:Resort)等整体布局相似、包含多个二义性建筑的场景容易出现误分现象。对于NWPU-RESISC45数据集,农田(9: circular_farm land)、港口(18:harbor)、雪山(38: snowberg)和存储罐(41: storage_ tank)等场景分类精度较高,和AID 数据集结果相似,商业区(11: commercial _area) 、教堂(8:church)和宫殿(28: palace)等整体布局相似且存在二义性地物的场景分类精度较低。

图6 训练比例为20%的AID 数据集生成的混淆矩阵Fig.6 Confusion matrix of AID dataset with training ratio of 20%

图7 训练比例为10%的NWPU-RESISC45 数据集生成的混淆矩阵Fig.7 Confusion matrix of NWPU-RESISC45 dataset with training ratio of 10%

为展示MAECO 的特征提取能力,绘制了MAE、DINO、MAECO 的注意力图以进行对比,如图9 所示。基于掩码重建任务的MAE 模型拥有良好的全局特征提取能力,但由于其在训练过程中对图像块进行高比例随机遮盖,造成了局部纹理等特征信息的丢失;基于对比学习的DINO 模型拥有良好的局部特征提取能力,但因其在知识蒸馏过程中不断舍弃次要知识,忽视了对次要地物的特征提取,如图8 中忽视了桥下的河流,可能导致对桥梁和公路误分;结合掩码重建任务和对比学习的MAECO 模型在提取全局特征的同时能够保留丰富的局部特征信息。

图8 注意力图对比Fig.8 Comparison of attention map

图9 MAE 和MAECO 的热力图对比Fig.9 Comparison of heat map between MAE and MAECO

2.5 消融实验

为了探究进行二阶段对比学习预训练对模型精度的影响,基于20%训练比例的AID 数据集,进行消融实验。实验设置分为3 组:1) 对于MAE 方法,设置掩码重建预训练阶段的迭代次数为200次;2) 对于结合单组对比学习的MAE 方法,于两个预训练阶段分别进行100 次迭代;3) 对于MAECO,在掩码重建预训练阶段和对比学习预训练阶段分别进行100 次迭代。上述实验分别在下游任务上迭代100 轮次,以控制每种设置下总迭代次数相当,结果如表5 所示。

表5 AID 数据集(训练比例20%)上的消融实验Tab.5 Ablation experiments on the AID detaset(20% of training)

由表5 可知,在MAE 基础上引入对比学习使精度提升了4.53%,相比MAE,使用对比学习模块的MAECO 的精度提升了6.22%。这是由于进行第二阶段的对比学习预训练可以使模型有效利用图像的局部特征,解决了掩码重建训练过程中局部特征丢失的问题,使得分类精度提高。

为了进一步直观展示对比学习模块的作用,使用Grad-CAM 方法对实验结果进行特征可视化。特征可视化生成的热力图如图9 所示。图中区域亮度越高表示模型对其关注度越高。由图9 可知,MAE 的关注区域较为离散,漏检、误检区域较多,对语义信息利用效率低下,这是由于MAE 存在局部特征信息丢失的问题,无法准确定位显著区域。MAECO 相比MAE 能够更加精准地定位并关注显著区域,这是由于MAECO 有效利用了局部上下文语义信息,以捕捉局部显著性特征,这充分证明了对比学习模块的有效性。

3 结论

针对MAE 模型训练过程中图像局部特征丢失的问题,本文提出了一种基于对比自监督学习的遥感图像分类模型MAECO。首先,使用掩码重建任务进行第一阶段预训练,提取图像全局特征;然后,使用对比学习模块进行第二阶段预训练,充分学习遥感图像丰富的局部语义信息,有效地补充了在第一阶段中丢失的局部特征;最后,构造线性分类器,使用少量有标注数据进行微调。在AID 数据集上和NWPU-RESISC45 数据集上进行对比实验,结果表明MAECO 模型精度优于SGSAGANs、 MoCo 等主流自监督模型, 接近ACNet、MG-CAP(Sqrt-E)等主流监督模型。可视化实验结果表明MAECO 模型能够有效捕捉全局和局部特征信息,具有良好特征提取能力。

目前本文在设计网络模型时,进行二阶段预训练,延长了模型拟合时间,下一步将尝试使用TransMix、自蒸馏方法进行模型加速,设计并行架构融合两个预训练阶段。