【AI气象2】AA-TransUNet: Attention Augmented TransUNetFor Nowcasting Tasks

目录

0. 论文基本信息

1. 背景介绍

2. 方法论

3. 实验

A. 提出的模型

B. 训练

C. 模型评估

4.实验

A. 降水图短临预测

B. 云覆盖预测

5.结果&讨论

6.总结


0. 论文基本信息

1. Paper地址:https://arxiv.org/abs/2202.04996

2. Code地址:https://github.com/YangYimin98/AA-TransUNet/tree/main/Models

3. [引用] Yang Y, Mehrkanoon S. Aa-transunet: Attention augmented transunet for nowcasting tasks[C]//2022 International Joint Conference on Neural Networks (IJCNN). IEEE, 2022: 01-08.

4. 摘要—近年来,基于数据驱动建模的方法在许多具有挑战性的气象应用中引起了广泛关注,包括天气要素预测。本文介绍了一种基于TransUNet的新型数据驱动预测模型,用于降水短时预报任务。TransUNet模型将Transformer和U-Net模型结合在一起,在医学分割任务中已经成功应用。在这里,TransUNet被用作核心模型,并进一步配备了卷积块注意模块(CBAM)和深度可分离卷积(DSC)。提出的Attention Augmented TransUNet(AA-TransUNet)模型在两个不同的数据集上进行了评估:荷兰降水图数据集和法国云覆盖数据集。所得结果显示,提出的模型在两个测试数据集上均优于其他受查模型。此外,提供了对提出的AA-TransUNet的不确定性分析,以提供有关其预测的额外见解。

5. 关键词—UNet、Transformer、降水短时预报、云覆盖短时预报、深度学习

6. 2022 International Joint Conference on Neural Networks (IJCNN)(CCF-C类)

1. 背景介绍

准确的天气预测对人类生活产生重要影响,直接影响着包括经济、农业、商业、交通和物流等多个领域。特别是短临降水预测,即未来6小时的天气预测,由于在支持上述工业部门的社会经济需求中具有高度影响,因此成为越来越受欢迎的研究主题[1]。此外,降水短临预测在水相关风险管理中起着重要作用。警告系统通常也依赖于此类预测来生成警报[2]。

传统的降水预测方法基于数值天气预报(NWP)[3]。然而,它们的空间和时间分辨率通常对降水短时预测来说太低。此外,基于NWP的模型需要大量的计算能力,因为它们通过模拟大气和海洋的基本物理过程来生成预测[4],[5]。因此,近年来出现了新的数据驱动模型,以克服NWP模型的缺陷。它们通常不依赖于大气和海洋的基本物理过程,而是专注于从历史数据中学习[6]。特别是,具有深度神经网络架构的模型在多个应用领域中表现出色[4],[6]–[15]。

基于卷积神经网络(CNN)的模型,如AlexNet [16]、ResNet [17]、卷积自动编码器(CAEs)[18]及其变种,是计算机视觉任务中最常用的神经网络架构之一。在CAEs的思想指导下,[19]中的作者引入了著名的UNet模型,该模型使用了经典的编码器(下采样)和解码器(上采样)结构,带有跳跃连接[19]。UNet最初用于医学图像分割,但逐渐在许多其他计算机视觉任务中取得了成功。

另一个具有影响力的模型是Transformer架构[20],它彻底改变了自然语言处理(NLP)领域。这种带有注意机制的序列到序列架构旨在处理序列输入数据。与用于序列到序列学习的传统模型相比,Transformer取得了更好的结果[20]。由于Transformer模型主导了NLP领域,许多研究人员开始探索其在计算机视觉任务中的适用性。例如,Google研究人员发布了Vision Transformer模型[21],该模型将图像分成固定大小的小块,线性嵌入每个小块,添加位置嵌入,并将生成的向量序列输入标准Transformer编码器。

尽管基于Transformer的模型在计算机视觉领域具有巨大潜力,但它们也面临着一些重大挑战,如对数据的高需求[22]。同样,一直占据着图像分割领域主导地位的UNet结构也存在一些限制,比如无法捕捉长期依赖性[23]。因此,将Transformer和UNet模型相结合已成为一种新的研究方向。基于这一点,[23]中的作者提出了TransUNet,这是一种图像分割模型,充分利用了UNet和Transformer的优势。TransUNet模型在医学图像分割方面表现出色,但在其他领域尚未广泛使用。

在本文中,我们提出了一种名为Attention Augmented TransUNet(AA-TransUNet)的新模型。在这里,TransUNet被用作核心模型,并进一步配备了Convolutional Block Attention Module(CBAM)[24]和Depthwise-separable Convolutions(DSC)[25]。通过在核心TransUNet模型中引入这些新元素,显著减少了可训练参数的数量,而不会降低模型的准确性。我们旨在将提出的AA-TransUNet应用于降水短时预测任务。特别是,在这里,我们使用与[5]中使用的相同数据集,即降水地图数据集和云覆盖数据集。与[5]中的情况类似,对于降水地图数据集,模型输入包括从地图雷达收集的60分钟降水地图,模型输出是对未来30分钟降水地图的另一个序列的预测。我们还在云覆盖数据集上进行了相关实验,以展示提出的模型的适用性。

本文的组织结构如下。第二节简要介绍了相关研究工作。第三节介绍了提出的AA-TransUNet模型。第四节提供了实验设置和所使用数据集的描述。第五节讨论了所得结果。最后,第六节得出了结论。

2. 相关工作

基于数据驱动的模型在气象要素预测,如温度、风速和降水等方面引起了广泛关注。气象要素预测可以被视为序列预测问题,因此与循环神经网络(RNN)模型,如长短时记忆(LSTM)[26],紧密相关。此外,由于天气数据的空间结构存在,它还与卷积神经网络(CNN)相关。因此,近年来出现了基于LSTM、CNN和UNet的强大深度学习架构,用于解决气象要素预测中的挑战性任务。

例如,[27]中的作者提出了卷积LSTM(ConvLSTM)模型,更好地捕捉了数据的潜在时空相关性,并始终优于普通LSTM。延续ConvLSTM研究方向,[28]中的作者提出了Convcast,一种基于嵌入式卷积LSTM的架构。[4]中的作者引入了不同的CNN架构,包括1-D、2-D和3-D卷积,以准确预测未来6到12小时的风速。UNet架构用于降水短时预测[29]。SmaAt-UNet模型是UNet模型的扩展,大大减少了UNet的参数,而不会影响其性能[5]。[6]中的作者引入了Broad-UNet,通过为UNet模型配备不对称并行卷积以及Atrous Spatial Pyramid Pooling(ASPP)模块来实现。

尽管UNet具有许多优点,但由于卷积操作的固有局部性,当建模长程依赖性时仍然可能具有有限的能力。另一方面,Transformer模型最初是为序列到序列预测开发的,使用全局自注意机制[20],但由于低级细节不足,也在local化能力上存在限制[23]。随后,[23]中的作者提出了TransUNet模型,该模型兼具UNet和Transformer模型的优点,并展示了它在医学图像分割任务中的成功应用。

在本文中,我们将TransUNet作为核心模型,并通过将Convolutional Block Attention Modules(CBAM)纳入核心模型,并在核心模型中使用Depthwise-separable Convolutions(DSC)进行扩展。CBAM是用于前馈卷积神经网络的轻量级通用注意模块。给定一个中间特征图,CBAM依次推断出两个维度的注意力图(通道注意模块和空间注意模块)[24]。深度可分离卷积由深度卷积(DW)和逐点卷积(PW)组成。这个结构类似于常规卷积操作,并可用于提取特征,尽管其参数大小远小于标准卷积操作[25]。

3. 方法论

本节介绍了提出的AA-TransUNet模型,该模型使用TransUNet [23]作为核心模型,并对其进行扩展以减少参数并提高预测性能。然后,我们研究了提出的模型在降水短时预测任务中的应用。

A. 提出的模型

本研究基于[23]的工作,扩展了TransUNet的概念。原始的TransUNet是一个编码器-解码器架构,其中使用CNN-Transformer作为编码器,同时使用原始的UNet解码器。这一模型已成功应用于医学图像分割任务。在这里,我们提出了一个新的编码器-解码器架构,称为AA-TransUNet,它将经典的TransUNet与两个关键元素相结合:卷积块注意模块(CBAM)和深度可分离卷积(DSC)。该模型如图1所示,并用于降水预测任务。我们将CBAM集成到TransUNet模型的编码器路径和解码器路径中。特别是,在编码器路径中,CBAM被整合到混合CNN-Transformer层的CNN部分中,在解码器路径中,它被放置在所有卷积层之后(参见图1)。这使得模型能够执行通道和空间上的注意力操作,因此可以更好地探索特征之间的通道关系和空间关系。此外,TransUNet解码器的原始卷积被DSC替换,旨在减小其参数大小。在表I中,我们比较了TransUNet和AA-TransUNet模型解码器的参数数量。尽管添加了CBAM层,但AA-TransUNet解码器的参数数量仍然比TransUNet低68.94%。需要注意的是,这里我们仅比较了解码器的参数数量。由于CBAM是一个轻量级通用模块[24],仅添加到编码器中,因此AA-TransUNet的编码器路径的参数数量与TransUNet相对相同。

与[5]不同的是,在这里,编码器的输入是应用了注意机制的图像特征。此外,与[5]相反,在解码器的每个级别中使用CBAM(请参阅图1中的解码路径中的紫色箭头)来识别图像的通道和空间维度中的重要特征[5],[24]。

与TransUNet [23]类似,AA-TransUNet模型的编码器是基于Transformer的模型,通常用于处理序列数据。因此,当将其用作编码器时,首先需要将图像序列化为2D块,以将3D图像转换为2D块[23]。然后将获得的块馈送到Transformer层,该层包括多头自注意力(MSA)层[20]和多层感知器(MLP)块(每个MSA和MLP通过层归一化门)。与此同时,与[23]一样,我们不使用纯Transformer作为编码器,而是使用CNN-Transformer,以利用解码路径中的中间高分辨率CNN特征图,并提高模型性能。

对于提出的AA-TransUNet解码器,使用了类似UNet解码器的架构。如图1所示,AA-TransUNet解码器中使用了四种操作:上采样(绿色箭头)、特征串联(虚线矩形)、双卷积(青色箭头)以及CBAM(紫色箭头)。模型的最后一层是1 x 1卷积(橙色箭头),产生一个表示预测结果的单一特征图。

B. 训练

我们将提出的AA-TransUNet模型与其他四个模型进行了比较,即持续性方法、UNet [19]、SmaAt-UNet [5]和TransUNet [23]模型,这些模型在两个数据集上进行了训练:降水图和云覆盖数据集。持续性方法是一种常用于短时预测任务的方法,它指的是使用序列的最后一幅图像作为预测图像的模型。其他四个基于UNet的模型在相同的设置下进行了训练。降水图数据集和云覆盖数据集的最大迭代次数分别设置为200和100。按照[5]的方法,采用了早停止准则方法,以提高训练效率并避免过拟合。在这种情况下,如果验证集在最后20个周期内没有增加,模型将停止训练。为了初始化学习率,测试了一系列值,即0.01、0.005、0.001、0.0005和0.0001,并经验性地发现0.001是最佳选择作为初始学习率。我们使用了一个学习率调度程序,在验证损失在四个周期内不增加时,将学习率减小到前一个学习率的十分之一。采用了Adam优化器,批大小设置为6。均方误差(MSE)被用作损失函数管尝试了不同数量的Transf。尽ormer层,即1、3、6、12和18层,但在所有进行的实验中,TransUNet和AA-TransUNet的Transformer层数量均设置为1,因为经验发现这是最佳选择。所有比较的模型都是在具有16GB RAM的NVIDIA Tesla P100上使用Pytorch进行训练的。模型的Pytorch实现可在GitHub上找到[1]。

C. 模型评估

为了评估模型的性能,报告了包括均方误差(MSE)、准确性、精确度、召回率和F1分数在内的五个指标的结果。与[5]类似,为了计算真正例(TP)、假正例(FP)、真负例(TN)和假负例(FN),模型输出和目标图像的每个像素都通过与固定阈值比较来转换为布尔掩码。阈值值在所有测试中都设置为0.5。

4.实验

在这一部分,给出了在两个数据集(降水图和云覆盖数据集)上进行的实验的详细信息,这两个数据集与[5]中使用的数据集相同。

A. 降水图短临预测

数据来自荷兰皇家气象研究所(Koninklijk Nederlands Meteorologisch Instituut,KNMI),先前已在[5]中使用过。该数据涵盖了2016年至2019年间荷兰王国及其邻国的降水信息。数据使用位于De Bilt和Den Helder的两个C波段多普勒天气雷达站收集。卫星每五分钟收集一次降水数据,并将其以图像格式存储。降水图的总数为420,000张。每个原始图像的尺寸为756 x 700,每个像素的值表示过去五分钟内每平方千米收集到的降水总量,并以整数格式存储。数据集的详细信息在[5]中提供,数据集可以在GitHub上找到[2]。

在这里也应用了与[5]中相同的预处理步骤。降水数据集中的所有图像都被裁剪为288 x 288,以便模型可以更专注于包含雨点的原始图像。此外,有很多图像中没有雨点,雨点像素和非雨点像素之间的不平衡较大。因此,与[5]一样,我们通过过滤原始降水数据集并仅选择至少包含50%和20%的像素含有任何降雨量的图像,创建了两个数据集,即NL-50和NL-20。裁剪前后数据集的信息可以在表II中找到。图2显示了NL-50数据集在30分钟内的示例。

作为预处理步骤,我们通过将它们除以NL-50训练集的最高值来对NL-50和NL-20数据集进行了归一化。每个数据样本是一个包含18个图像的序列。为了进行公平比较,与[5]类似,我们使用前60分钟的数据(前12个图像)作为模型的输入,将最后30分钟的数据用于与模型预测进行比较。此外,使用从2016年到2018年的数据集作为训练集,使用2019年的数据集作为测试集。将随机选择10%的训练集作为验证集,以提示何时停止训练。

NL-50的训练集用于训练模型。NL-50的验证集用于模型选择。一旦所有模型都训练完成,它们将在NL-50测试集和NL-20数据集上进行测试。通过在NL-20测试集上测试模型,可以检查模型的泛化能力。

B. 云覆盖预测

云覆盖数据集在[18]中引入,并先前在[5]中使用过。云覆盖数据集根据云的高度和类型被分成16类。此分类是通过Meteosat Second Generation(MSG)拍摄的各种可见光和红外线通道图像计算的,MSG是一颗位于0度经度的静止卫星,于2017年和2018年在法国使用。卫星每15分钟收集一次数据,包含3712 x 3712大小的图像。

在云覆盖图像中,每个像素可能有15种不同的值。云覆盖状态的详细信息在表III中报告。数据根据图像中是否有云覆盖而进行二值化处理。更具体地说,与[5]一样,将值1到4转换为0(无云),将值5到15转换为1(有云)。随后,根据法国的边界,将图像裁剪为256 x 256的大小[6]。

图3显示了云覆盖数据样本。与[5]类似,我们基于前60分钟(前4张图像)来预测过去90分钟(最后6张图像)的云覆盖状态。与[5]一样,我们使用2017年和2018年上半年的数据作为训练集,并随机选择其中的10%作为验证集。然后,2018年数据集的其余部分用作测试集。与先前的降水数据集处理方法不同,云覆盖数据集不进行归一化处理,因为它只包含具有二进制值0和1的图像。

5.结果&讨论

在这一部分,给出了对三个数据集(NL-50、NL-20和云覆盖)上研究模型的测试结果以及相应的分析。

A. 降水图短时预测评估

表IV显示了在NL-50测试集上获得的结果。从表IV中可以看出,所有模型都优于持续基线模型。需要注意的是,由于短时间内天气数据变化不大(降水预测和云覆盖预测的输入仅为1小时),持续模型已经能够提供良好的预测结果。然而,表IV显示,在NL-50测试集上,即使表现最差的基于UNet的模型,即TransUNet,其均方误差(MSE)值也不到持续模型的一半。

提出的AA-TransUNet模型在所有研究模型中获得了最低的MSE值。TransUNet模型的MSE值为0.0124,接近UNet和SmaAt-UNet的MSE值。然而,在加入卷积块注意模块(CBAM)和深度可分离卷积(DSC)后,MSE得分降至0.0110,表明所提出的模型比传统的TransUNet提高了11%。此外,需要注意的是,在TransUNet解码器中将普通卷积替换为DSC后,TransUNet解码器的参数数量大大减少。此外,除了获得的MSE值外,所提出的模型在其他指标上也表现良好。

图4显示了NL-50和NL-20的地面实况样本图像以及AA-TransUNet的预测示例。可以看出,NL-50的样本图像比NL-20有更多的雨点像素。通过比较地面实况和模型预测,可以观察到我们提出的模型可以准确预测雨点像素的集中位置。与地面实况相比,模型的预测结果略微模糊。这是因为我们在训练模型时使用了均方误差(MSE)作为损失函数,而MSE损失倾向于朝着产生模糊图像的梯度下降的方向执行[30]。

我们还在NL-20测试集上测试了模型。表V显示,尽管NL-20上的雨点像素密度远低于NL-50,我们提出的AA-TransUNet模型仍然优于其他模型。表V显示,所有模型的预测结果都远远超过了持续模型的结果。此外,AA-TransUNet的性能,它获得了最低的MSE值,也比传统的TransUNet提高了5%。

此外,我们还对NL-20的预测图与地面实况图进行了比较。可以从图4(c, d)看出,尽管地面实况雨点像素的密度非常低,但我们的模型仍然几乎预测出了所有雨点像素的轮廓。

接下来,我们分析了不同数量的Transformer层对AA-TransUNet模型在NL-50测试集上性能的影响。如图1所示,每个Transformer层包括多头自注意力(MSA)和多层感知器(MLP)块。我们已经检查了有1、3、6、12和18个块(对应于1、3、6、12和18个Transformer层)的情况。根据图5,具有一个Transformer层(一个MSA和一个MLP块)的模型在NL-50验证集上表现最佳。这一观察结果与[31]–[33]的研究结果一致,即当数据集大小较小时,使用较少数量的块(Transformer层)的Transformer可以获得更好的性能。

B. 云层预测的评估

上述模型在云层数据集上进行了训练和测试,获得的结果在表VI中列出。从表VI可以观察到与降水图数据集相同的性能模式。所有模型都取得了远远超过持续模型的结果。同时,AA-TransUNet在所有模型中具有最低的MSE值以及最高的准确性和精度值。值得一提的是,云层数据集上所有模型的MSE值都高于之前在降水数据集上获得的值,因为云层数据集只包含二进制值。因此,实际标签的值有时与预测值非常不同,因此导致更高的MSE值。

图6显示了AA-TransUNet的预测和云层数据的地面实况图像的示例。该模型已准确识别了256 x 256大小图像上的所有云层的出现。

C. 不确定性分析

我们通过Epistemic不确定性来量化不确定性,它衡量了模型预测的不确定性。理论上,Epistemic不确定性可以通过贝叶斯神经网络[34]来获取。然而,与[34]、[35]一样,我们在这里使用测试时的dropout(TTD)来近似贝叶斯推断。具体而言,在测试时,对于每个目标降水图,通过运行模型k次(在我们的情况下,我们设置k=20)来生成后验分布。然后,通过取得到的k个降水图(后验采样分布)的方差来测量模型的不确定性图。此外,我们还取后验采样分布的均值,它可以被视为模型的预测,并被称为TTD预测。图6和图7显示了云层以及NL-20数据集的获得的TTD不确定性图和TTD预测。此外,图8显示,随着预测步数的增加,提出的AA-TransUNet模型的不确定性也会增加。

6.总结

本文提出了一种新颖的Attention Augmented TransUNet模型(AA-TransUNet),用于降水现在预测任务。我们展示了将卷积块关注模块(CBAM)和深度可分离卷积(DSC)纳入经典TransUNet模型中,可以提高性能,同时显著减少模型解码器部分的参数数量。我们展示了所提出模型在两个现在预测任务上的适用性,即降水和云层现在预测。实验结果表明,AA-TransUNet模型优于原始TransUNet模型和其他检验的模型。我们提出模型的实现,包括训练好的模型,可在Github上找到:https://github.com/YangYimin98/AA-TransUNet。

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值