全球首个面向遥感任务设计的亿级视觉Transformer大模型

语言: CN / TW / HK
作者:京东探索研究院

深度学习在很大程度上影响了遥感影像分析领域的研究。然而,大多数现有的遥感深度模型都是用ImageNet预训练权重初始化的,其中自然图像不可避免地与航拍图像相比存在较大的域差距,这可能会限制下游遥感场景任务上的微调性能。

 

为此,京东探索研究院联合武汉大学、悉尼大学借助迄今为止最大的遥感场景标注数据集MillionAID,从头开始训练包括卷积神经网络(CNN)和已经在自然图像计算机视觉任务中表现出了良好性能的视觉Transformer(Vision Transformer)网络,首次获得了一系列基于监督学习的遥感预训练基础骨干模型。并进一步研究了ImageNet预训练(IMP)和遥感预训练(RSP)对包括语义分割、目标检测在内的一系列下游任务的影响。

 

实验结果证实了探索研究院先前提出的先进Vision Transformer系列模型ViTAE在遥感任务上的优越性,并发现RSP在遥感任务上的有效性以及感知相关语义方面具有的独特性。实验结果进一步表明RSP会受到上下游任务差异的影响,这些发现对遥感大规模数据集和预训练方法提出了新的要求。

  image.png

01

研究背景

近年来,深度学习凭借自动提取反映物体固有属性的深度特征的优势,在计算机视觉领域取得了令人印象深刻的突破,遥感领域也不例外。在遥感领域,最常用的是深度模型是卷积神经网络(CNN)。 目前,几乎所有的遥感深度模型都是在计算机视觉领域最著名的图像数据集 ImageNet-1K 数据集上进行预训练,该数据集中来自 1,000 个不同类别的百万张真实世界图像使模型能够学习强大的表示。然后这些预训练后的模型被可以用做遥感任务的骨干网络进行进一步微调。

 

尽管这些模型在遥感任务中取得了显著的效果,但仍有一些问题需要研究。直观地说,与自然图像相比,遥感图像在视角、颜色、纹理、布局、对象等方面明显存在较大的域差距。以前的方法试图通过进一步微调遥感图像数据集上的预训练模型来缩小这一差距。然而,ImageNet预训练(IMP)引入的系统偏差对性能有着不可忽视的副作用。另一方面,我们注意到,随着遥感技术的进步,各种各样的传感器捕捉到了丰富的遥感图像,可以用于预训练。作为一个代表性的例子,MillionAID 是迄今为止最大的遥感图像数据集,它是从包含多种传感器图像的谷歌地球(GE)上收集的,并且具有类似ImageNet-1K的百万级图像数量规模,这使遥感预训练(RSP)成为可能。

 

RSP能够从头开始训练深度模型,这意味着候选模型不必局限于现成的CNN。因此,在本文中,我们也研究了视觉Transformer(Vision Transformer)的骨干网络,它们在计算机视觉领域表现出了令人惊讶的性能。与CNN中擅长局部建模的卷积相比,Vision Transformer中的多头自注意(MHSA)能够灵活地捕捉不同的全局上下文。最近,探索研究院提出的ViTAE模型探索了卷积和MHSA的平行结构,以同时建模局部性和长程依赖性,在ImageNet分类任务和下游视觉任务上取得了很好的结果。此外,它还通过扩张卷积模块和层级设计提取多尺度特征,这对于计算机视觉下游任务,尤其是在遥感图像理解任务,都具有重要的价值。因此我们研究了CNN和层级Vision Transformer网络经过RSP后,在场景识别、语义分割、目标检测和变化检测等遥感任务上的微调性能。为了实现这些目标,我们在九个流行的数据集上进行了广泛的实验,并得出了一些有益的结论。RSP是遥感图像理解中的一个新兴研究方向,但仍处于探索阶段,尤其是基于Vision Transformer这种新型网络架构的预训练方法。我们希望这项研究能够填补这一空白,并为未来的研究提供有用的见解。

 

02

MillionAID, ViTAE 和ViTAEv2的介绍

 

1.MillionAID

MillionAID 是迄今为止遥感领域最大的数据集。它包含 100,0848 个不重叠的场景,有51类,每类大约有2,000-45,000图像。该数据集来自谷歌地球,由包括但不限于 SPOT、IKONOS、WorldView 和 Landsat 系列的多种传感器组成,因而图像分辨率不同。最大分辨率可达0.5m,最小的则有153m。图像大小范围从 110110 到 31,67231,672。该数据集均为RGB图像,非常适合训练典型的视觉神经网络模型。

 

2.ViTAE 和ViTAEv2

   image.png

ViTAE是探索研究院最近提出的先进Vision Transformer模型,它采用深窄式设计,在网络开始时迅速降采样,然后将网络加深,在提高性能的同时降低模型大小和计算成本。ViTAE 模型首先通过三个Reduction Cell将输入图像下采样到1/16分辨率。与 ViT 类似,在添加位置编码之前,将class token与第三个Reduction Cell的输出连接。然后堆叠多个Normal Cell,并始终保持特征图分辨率。最后一个Normal Cell的class token输入到线性层进行分类。ViTAE 模型在 ImageNet 数据集上分类性能表现出色,但它不方便像CNN那样产生层次化的中间特征,从而迁移到分割、检测和姿态估计等其它下游任务(目前有一些新技术来解决这个问题,例如ViTDet,并已取得较好成果,请关注我们的复现Repo:https://github.com/ViTAE-Transformer/ViTDet)。

 

在此基础上,探索研究院提出了 ViTAEv2,它采用了 ResNet 和 Swin 等流行骨干网络的层次化设计。在 ViTAEv2 中,网络被分成四个阶段。每个阶段首先采用Reduction Cell进行下采样,然后堆叠多个 Normal Cell进行特征变换。在最后一个阶段的Normal Cell后使用全局平均池化层来替换class token。当对下游任务进行微调时,该池化层被移除,剩下的网络与相应任务的解码器相连。图 2 显示了原始 ViTAE 和 ViTAEv2的网络架构。

  image.png

Reduction Cell和Normal Cell是 ViTAE 中最重要的两个模块,它们是基于典型的Transformer模块来构建的。Reduction Cell用于下采样并提供多尺度上下文。具体来说,在输入归一化层和 MHSA 层之前,这些特征会通过一个金字塔缩减模块(PRM)。该模块包含具有不同扩张率的多个并行的扩张卷积,其中步幅大小控制了空间降采样率。在 PRM 后,来自平行分支的特征在通道维连接。PRM 将 CNN 的尺度不变性引入 ViTAE,而局部性建模则是通过将输入到PRM 的特征同时送入到并行卷积模块 (PCM) 中来完成。PCM 位于与包含 PRM 和 MHSA 的全局依赖路径平行的附加分支中,它由三个连续的卷积层组成。通过调整步幅,PCM 的下采样率与 PRM 相同。来自MHSA、PCM 和原始残差分支的三个特征在输入前馈网络(FFN)之前进行加法融合。需要注意的是,Normal Cell和Reduction Cell具有类似结构,但是不包括 PRM模块。

 

受 Swin Transformer的启发,ViTAEv2中上述cell中的一些 MHSA 被替换为窗口MHSA(WMHSA)以降低计算成本。考虑到后期特征尺寸变小,不需要用窗口划分特征。因此,只有前两个阶段的 MHSA 被 WMHSA 替代。需要说明的是,ViTAEv2采用的 WMHSA 不需要像Swin Transformer那样进行循环偏移,因为 WMHSA 是在 PRM 的合并多尺度特征上进行的,其中不同区域之间已经通过扩张卷积的重叠感受野实现了信息交换。此外,因为卷积已经能够编码位置信息,ViTAEv2也不需要再使用相对位置编码。ViTAE 和 ViTAEv2 中不同cell的详细结构和比较如图3所示。

 

在本次研究中,我们主要评估原始 ViTAE 的“Small”版本,名为 ViTAE-S。相应的,我们还采用了 ViTAEv2-S 模型,因为它具有出色的表征能力和对下游任务更好的可迁移性。

 

03

遥感预训练的实施

 

1. 确定预训练模型

我们首先确定用于RSP的深度模型的类型。为此,我们从MillionAID官方训练集中构建了一个迷你训练集和迷你评估集,分别有9775和225张图像。注:后一组是通过从每个类别中随机选择5张图像来平衡类别。对于CNN,使用了经典的ResNet-50 。由于本研究主要探讨RSP下的CNN和Vision Transformer模型的性能,因此我们还评估了一系列典型的基于Vision Transformer的网络,包括DeiT-S 、PVT-S 和Swin-T。选择特定版本模型的一个考量是为了确保这些模型和ResNet-50以及ViTAE-S模型具有相似的参数量。此外,考虑到ViT是视觉Transformer的最基本模型,我们选择了其最小版本ViT-B 模型以供参考。

image.png

image.png

表II展示了各个模型的结果,可以看出,尽管ViT-B的参数最多,但其性能不如经典的ResNet-50。DeiT-S表现最差,因为我们没有采用教师模型辅助训练。由于我们的任务是利用遥感图像进行预训练,因此获取相应的教师模型可以认为是我们的目标而不是前提。通过引入特征金字塔的设计范式,PVT-S与ViT-B相比提高了准确性。在此基础上,原始ViTAE-S模型进一步考虑了局部性和尺度不变性这些传统CNN具有的归纳偏置。

 

然而,由于早期下采样模块(Reduction Cell, RC)中的特征分辨率较大,需要更多的计算,因此需要花费较多的训练时间。Swin-T通过在固定窗口中限制MHSA来解决这个问题,并采用窗口偏移来隐式促进窗口之间的通信。ViTAEv2引入了这种窗口多头自注意力(Window MHSA, WMHSA),并因为卷积旁路已经能够促进跨窗信息交互,从而省去了窗口偏移和相对位置编码操作。最终,ViTAEv2-S实现了最佳性能,并以2.3%的top-1准确率超过了第二名。

 

基于上述结果,我们选择候选模型的具体程序如下。首先,我们选择ResNet-50作为常规CNN中的代表网络。经过遥感预训练的ResNet-50,可以在一系列遥感数据集上提供一组新的CNN参考基线。由于准确率低、参数多,我们没有选择DeiT-S和ViT-B模型作为候选模型。此外,由于堆叠Transformer的设计,它们很难迁移到下游任务中。(目前有一些新技术来解决这个问题,例如ViTDet,并已取得较好成果,请关注我们的复现Repo:https://github.com/ViTAE-Transformer/ViTDet)。

Swin Transformer也具有PVT的特征金字塔结构,并采用WMHSA取代全局MHSA,节省了显存和计算量。由于Swin-T的top-1精度大于PVT且需要的训练时间较少,因此我们在后续实验中也选择了Swin-T作为候选模型。对于ViTAE模型,我们选择性能最强的模型,即ViTAEv2-S,以期望在后续任务(如遥感场景识别)中具有良好的性能。

 

2. 获得合适的权重

在确定上述候选模型后,我们对它们进行 RSP 以获得预训练的权重。具体来说,为了保持类别平衡,我们在 MillionAID 数据集的每个类别中随机选择 1,000 张图像,形成包含 51,000 张图像的验证集,与包含 50,000 张图像的 ImageNet 验证集的规模相当,并把剩下的 949,848 张图像用于训练。 

image.png

为了获得合适的预训练权重,我们在不同训练代数(epoch)的配置下分别训练 ViTAEv2-S 模型。结果如表III所示。可以观察到模型在大约 40 个 epoch 后开始性能饱和,因为与训练 20 个 epoch 相比,top-1 准确率仅提高了 0.64%,而接下来的 20 个 epoch 只带来了 0.23% 的增益。因此,我们首先选择训练了 40 个 epoch 的网络权重作为 ViTAEv2-S 的 RSP 参数,并应用于后续任务。直觉上,在大规模预训练数据集上表现良好的模型在下游任务上也会表现良好。因此,我们还在下游任务中使用了经过 100 个 epoch 训练的网络权重。这些模型分别用后缀“E40”和“E100”表示。

 

对于 ResNet-50 和 Swin-T,我们遵循Swin 的训练设置,即模型训练了 300 个 epoch。在实验中,我们观察到 Swin-T-E120 在验证集上的 top-1 准确率大致相当于 ViTAEv2-S-E40。因此,我们也选择了 Swin-T-E120 的训练权重。同样,我们也选择了最终的网络权重 Swin-T-E300 作为与 ViTAEv2-S-E100 的比较。为了使实验公平,还考虑了使用 40 个 epoch 训练的 ResNet-50 和 Swin-T 的权重,因为它们与 ViTAEv2-S-E40 经过了同样的训练代数。

  image.png

最终的预训练模型列在表IV中。可以看出,验证集准确率几乎随着训练 epoch 的增加而增加。但是,Swin-T-E300 的性能略低于 Swin-T-E120。尽管如此,我们仍然保留了Swin-T-E300 模型。因为模型在训练阶段见到更多的样本之后,它可能具有更强的泛化能力。

 

04

下游任务上的微调实验

1. 场景识别   image.png

定量实验: 表 V 展示了使用不同方法预训练的上述候选模型和其他 SOTA 方法的结果。最后三组中的粗体字表示每组中最好的结果,而“*”表示所有模型中最好的(在其他任务中含义相同)。与 ImageNet 预训练的 ResNet-50 相比,我们的遥感预训练 ResNet-50 在所有设置下均提高了准确性。这些结果意味着 RSP 为后续微调过程的优化带来了更好的起点。同样,RSP-Swin-T 在三个设置上的表现优于 IMP-Swin-T,在其他两个设置上也取得了可比较的结果。此外,与其他复杂方法相比,ResNet-50 和 Swin-T 仅使用 RSP 权重而不改变网络结构的情况下就取得了有竞争力的结果,从而证明了遥感预训练的价值。

 

此外,在比较 ImageNet 预训练的 ResNet-50 和 Swin-T 时,我们可以发现 IMP-Swin-T 在所有设置上的表现都更好,因为Vision Transformer具有更强的上下文建模能力。不过在通过 RSP 权重进行初始化后,ResNet 变得更具竞争力。由于 ViTAEv2-S同时具有局部建模能力和远程依赖建模能力,无论 IMP 和 RSP,它在几乎所有设置上都优于 ResNet-50 和 Swin-T。此外,RSP-ViTAEv2-S 在除 AID (5:5) 之外的几乎所有设置上都实现了最佳性能。

  image.png

定性实验:图4 显示了不同评估模型来自各种场景的图像的不同区域的响应。与 IMP-ResNet-50 相比,RSP-ResNet-50 更关注重要目标。这意味着 RSP 有助于 ResNet-50 学习更好的表示,归功于 MillionAID 数据集中提供的大量语义相似的遥感图像。令人惊讶的是,IMP-Swin-T 模型主要关注背景区域,但经过 RSP 之后,其前景响应得到了显著增强。ViTAEv2-S通过结合CNN和视觉转换器的优势,同时具备局部和全局上下文捕捉能力,实现了对整个场景的全面感知。RSP-ViTAEv2-S 不仅关注主要对象,还考虑了背景中的相关区域。在前景物体上,RSP-ViTAEv2-S 也能给予更高的关注度,在对象分布复杂的场景,RSP-ViTAEv2-S能够形成统一且完整的地物表征,有效感知场景的整体信息。

 

2. 语义分割

image.png

定量实验: 表VII 展示了采用UperNet 框架时,我们的方法和其他 SOTA 方法在iSAID数据集上的分割结果。可以看出,将骨干网络从 ResNet-50 更改为 Swin-T,再更改为 ViTAEv2-S 时,性能有所提高。结果与上述场景识别结果一致,表明视觉Transformer具有更好的表示能力。另一方面,经过ImageNet预训练的IMP-Swin-T 取得了具有竞争力的结果,而IMP-ViTAEv2-S 在 iSAID 数据集上取得了最佳性能。表VII 还显示了 RSP 模型的优势在于感知一些具有明确遥感语义的类别,例如“桥梁”,这符合之前场景识别任务中的发现。

  image.png

定性实验: 图 6 中展示了在 Potsdam 数据集上采用不同预训练骨干网络的UperNet 分割模型的一些视觉分割结果。对于长条形地物,其长度较长,要求模型能够捕获长程上下文,而宽度又较窄,对模型的局部感知能力又提出了要求,而ViTAEv2网络因为将CNN的局部性和尺度不变性引入到Vision Transfomer网络中,同时具有了CNN和Transformer的优势,因而能够同时实现全局和局部感知。因此,只有ViTAEv2-S成功连接了长条状低矮植被(如红框所示)。

3. 目标检测

  image.png

定量实验: 表VIII 显示了 目标检测实验的结果。在具有挑战性的 DOTA 数据集上,可以看出使用先进的 ORCN 检测框架,采用ResNet-50 或 Swin-T 骨干网络的模型表现良好。ViTAEv2-S 通过引入了 CNN 的局部性和尺度不变性等归纳偏差,获得了惊人的性能,将 ORCN 基线提高了近 2% mAP。需要注意的另一点是,RSP在这三个骨干网络上的性能都优于IMP。RSP-ViTAEv2-S 的总体mAP比IMP-ViTAEv2-S 高,因为 RSP 在“桥梁”以及包括“直升机”和“飞机”在内的飞行器类别上具有显着优势,而在其他类别上,这两种模型之间的差距并不很大。

image.png

定性实验: 图7 可视化了 DOTA 测试集上使用 ViTAEv2-S 骨干网络的 ORCN 模型的一些检测结果。红框表示,当对象密集分布时,RSP-ViTAEv2-S 仍然可以预测正确的对象类别,而 IMP-ViTAEv2-S 被密集上下文混淆并做出错误预测。对于长条形的“桥梁”类别,IMP-ViTAEv2-S 产生了漏检(见黄色框),而 RSP-ViTAEv2-S 模型以更高的置信度分数成功检测到该物体,这再一次呼应了先前的发现。

4. 变化检测

  image.png

定量实验: 表X 中展示了采用不同预训练骨干网络的BIT 框架在变化检测任务上的定量实验结果。可以看到,自监督的 SeCo 预训练权重在此任务上表现良好,虽然SeCo 的目标是经过对比学习以实现季节不变性特征学习,但由于其采用了多头子空间嵌入的方式对变化特征进行编码,所以其仍然能在具体的分支上学习到对季节变化敏感的特征表示。尽管如此,通过 IMP 或 RSP 预训练的 ViTAEv2-S 的性能优于 SeCo-ResNet-50,显示了使用先进骨干网络的好处。与其他方法相比,ViTAEv2-S取得了最佳性能,显示了将先进的Vision Transformer模型应用于遥感领域的潜力。

 

通过不同模型在不同任务下RSP和IMP下的性能对比,我们可以推断出变化检测所需表示的粒度应该介于分割和检测之间,因为它虽然是一个分割任务,但是只有两个类别,不需要去识别特定的语义类别。

 

 

定性实验: 图8展示了一些视觉变化检测结果。可以看出,IMP 的 ResNet-50 和 Swin-T 并不能很好地检测到自然场景中田野内道路的变化。采用 RSP 可以部分缓解这个问题。SeCo-ResNet-50 进一步提高了道路区域的检测,这与表X 中的结果一致。与上述模型相比,ViTAEv2-S 模型有效地捕捉到了道路细节。在人工变化的场景中,ViTAEv2-S 模型解决了所有其他模型结果中存在的对象粘连问题,这表明 ViTAEv2-S 的特征在区分物体和背景方面更具判别力。

 

5. 不同遥感预训练骨干网络的综合比较

  image.png

最后,我们全面比较了 RSP 在所有任务上的不同骨干网络的性能。具体来说,我们对每个任务的所有数据集的分数进行平均计算,结果如表XI。可以发现,预训练更多 epoch 的骨干通常在下游任务上表现更好,因为它们获得了更强的表示。尽管也有例外,例如预训练300代的 Swin-T 模型在对象检测任务表现不如预训练120代的对应模型,这暗示任务差异也很重要。结合了 CNN 和Vision Transformer优势的ViTAEv2-S模型在所有任务上都表现出了最好的性能。

 

05

结论

 

在这项研究中,我们在最大的遥感数据集 MillionAID 上研究了基于 CNN 和Vision Transformer的遥感预训练问题,并综合评估了它们在场景识别、语义分割、对象检测和变化检测四个下游任务上的表现,并将它们与 ImageNet 预训练和其他 SOTA 方法进行比较。通过综合分析实验结果,我们得出以下结论:

(1) 与传统的 CNN 模型相比,视觉Transformer在一系列遥感任务上表现出色,特别是 ViTAEv2-S这种将 CNN 的固有归纳偏置引入到Vision Transformer的先进模型,在这些任务的几乎所有设置中都实现了最佳性能。

(2) 经典 IMP 使深度模型能够学习更通用的表示。因此,IMP 在处理遥感影像数据时,仍可以产生具有竞争力的基线结果。RSP 产生了可与 IMP相当或者更好的结果,并且由于减轻了上游预训练任务和下游任务之间的数据差异,因此在某些特定类别(例如“桥梁”和“飞机”)上表现更为出色。

(3) 任务之间的差异对 RSP 的性能也有影响。如果特定下游任务所需的表示更接近上游预训练任务(例如场景识别),则 RSP 通常会带来更好的性能。

 

我们希望这项研究可以为遥感社区提供有关使用先进Vision Transformer和遥感预训练的有用见解。为了方便大家使用,所有遥感预训练模型及相关代码均已开源,详见https://github.com/ViTAE-Transformer/ViTAE-Transformer-Remote-Sensing 。另外,关于采用非层次化Vision Transformer模型应用于下游任务的研究进展,可以关注ViTDet的方法以及我们的复现代码:https://github.com/ViTAE-Transformer/ViTDet 。我们也会在ViTAE-Transformer-Remote-Sensing的官方repo中不断更新相应的结果。

  论文链接:https://arxiv.org/abs/2204.02825

项目地址:https://github.com/ViTAE-Transformer/ViTAE-Transformer-Remote-Sensing

 

参考文献

[1] D.Wang, J. Zhang, B.Du, G-S.Xia and and D. Tao, “An Empirical Study of Remote Sensing Pretraining”, arXiv preprint, axXiv: 2204: 02825, 2022.

[2] Y. Long, G.-S. Xia, S. Li, W. Yang, M. Y. Yang, X. X. Zhu, L. Zhang, and D. Li, “On creating benchmark dataset for aerial image interpretation: Reviews, guidances and million-aid,” IEEE JSTARS, vol. 14, pp. 4205–4230, 2021.

[3] Y. Xu, Q. Zhang, J. Zhang, and D. Tao, “Vitae: Vision transformer advanced by exploring intrinsic inductive bias,” NeurIPS, vol. 34, 2021.

[4] Q. Zhang, Y. Xu, J. Zhang, and D. Tao, “Vitaev2: Vision transformer advanced by exploring inductive bias for image recognition and beyond,” arXiv preprint arXiv:2202.10108, 2022.

[5] T. Xiao, Y. Liu, B. Zhou, Y. Jiang, and J. Sun, “Unified perceptual parsing for scene understanding,” in ECCV, 2018, pp. 418–434.

[6] X. Xie, G. Cheng, J. Wang, X. Yao, and J. Han, “Oriented r-cnn for object detection,” in ICCV, October 2021, pp. 3520–3529.

[7] H. Chen, Z. Qi, and Z. Shi, “Remote Sensing Image Change Detection With Transformers,” IEEE TGRS., vol. 60, p.3095166, Jan. 2022.

[8] Y. Li, H. Mao, R. Girshick, K. He. Exploring Plain Vision Transformer Backbones for Object Detection[J]. arXiv preprint arXiv:2203.16527, 2022.