【AI气象3】SmaAt-UNet: Precipitation Nowcasting using a SmallAttention-UNet Architecture

目录

0. 论文基本信息

1. 背景介绍

2. 相关工作

3. 方法论

3.1. 提出的 SmaAt-UNet 模型

3.2. 其他模型

3.3. 训练

3.4. 模型评估

4.实验

4.1. 降水图数据集

4.2. 云量数据集

5.结果&讨论

5.1. 降雨图数据集上的评估

5.2. 云覆盖数据集上的评估

6.总结

---

0. 论文基本信息

1. Paper地址：https://arxiv.org/abs/2007.04417

2. Code地址：https://github.com/HansBambel/SmaAt-UNet/tree/master

3. [引用] Trebing K, Staǹczyk T, Mehrkanoon S. SmaAt-UNet: Precipitation nowcasting using a small attention-UNet architecture[J]. Pattern Recognition Letters, 2021, 145: 178-186.

4. 摘要—气象预测主要由数值天气预报所主导，它试图准确建模大气的物理特性。数值天气预报的一个不足之处在于它缺乏使用最新可用信息进行短期预报的能力。通过使用基于数据驱动的神经网络方法，我们展示了可以生成准确的降水现在预报。为此，我们提出了SmaAt-UNet，这是一种高效的卷积神经网络，基于众所周知的UNet架构，配备了注意力模块和深度可分离卷积。我们使用荷兰地区的降水图和法国云覆盖的二进制图像对我们的方法进行了评估。实验结果显示，就预测性能而言，所提出的模型与其他检验的模型相当，但只使用了可训练参数的四分之一。

5. 关键词—UNet、降水短临预报、深度学习、轻量化

6. Pattern Recognition Letters（SCI顶刊）

1. 背景介绍

计算机气象预测是现代工业化社会中普遍存在的特征，用于规划、组织和管理个人和经济生活的各个方面。迄今为止，天气预报的主要方法是数值天气预报（NWP）。NWP依赖于数学模型，考虑大气的不同物理特性，如风速、压力和温度。基于NWP的模型可以生成几小时到几天后的准确天气预测。然而，它们涉及求解高度复杂的数学模型，计算成本高昂，需要庞大的计算能力，因此通常在昂贵的超级计算机上执行[1]。

由于NWP模型具有高计算和时间要求，因此不太适合范围从几分钟到最多6小时的短期预报，也称为短临预报[2]。现在预报模型能够使用最新的可用观测天气数据来创建其预测，使其比NWP模型更具响应性[3]。这种响应性对于提高对动态和快速变化的环境（如大气层）的准确预测至关重要。因此，现在预测已经成为重要的工具，以补充NWP方法，特别是在典型的严重天气危害，如雷暴和大雨，常见的气象不稳定条件下[3]。正如美国气象学会的一份状态报告所强调的，现在预测雷暴在航空、建筑业、电力公用事业和地面交通等各个领域都有相关应用[4]。短临预测还在2008年北京奥运会上用于确保运动员的安全[5]。最后，天气短临预测还可以用于规划日常生活的普通活动。

最近，人工神经网络体系结构（ANNs）的进展使基于数据驱动的模型能够弥合短期预测的现有差距[6,7,8,9]。NWP和ANN之间的关键区别在于前者是模型驱动的，而后者是数据驱动的方法。与基于模型的方法不同，数据驱动的模型不依赖于大气底层物理的计算来进行预测。相反，它们分析并学习历史气象数据，如过去的风速和降水图，以预测未来。

在本文中，我们介绍了一种新颖的基于人工神经网络的模型，用于预测未来30分钟内的高分辨率格网降水。我们模型的输入数据包括降水图，即显示一段时间内累积降雨量的地图雷达图像。此外，我们还展示了该模型在云层现在预测任务上的适用性。

在先前的研究中，卷积神经网络被描述为处理图像数据的有效方法。卷积是基于核的操作，可以在图像上滑动，使模型能够以比其他前馈方法更高效的方式捕捉局部不变特征[10]。它们已经成功应用于各个领域，不仅包括图像处理，还包括其他类型的信号处理。例如，[11]的作者使用基于CNN的模型为输入图像创建标题，而[12]采用了CNN来检测图像中的对象。[13]的作者介绍了一个基于3D CNN的模型，用于预测丹麦不同城市的风速。在另一项研究中，应用了基于CNN的架构，对来自智能手机加速计的信号进行分类，以识别用户的交通方式[14]。[15]的作者引入了用于风速预测的多维卷积神经网络。

考虑到CNN在涉及图像输入的任务中的有用性，它们为降水现在预测提供了有前途的解决方案。在本文中，我们提出了Small Attention-UNet（SmaAt-UNet）模型。它以UNet架构作为核心模型，并配备了注意力模块和深度可分离卷积（DSC）。我们的模型有一个优点，即能够将模型参数的大小减小到原始UNet实现的四分之一，同时保持与原始UNet架构相当的性能。

模型尺寸的减小为在小型计算单元上使用降水模型提供了可能性，例如智能手机，类似于[17]。这可以通过在几秒钟内根据用户请求创建使用最新可用数据的预测来实现。此外，与更大的模型具有相似性能的模型尺寸减小对于创建需要更少训练和计算能力的高效架构至关重要。

本文的组织如下。第2节简要概述了使用机器学习架构进行天气预测的相关研究。在第3节中，我们描述了提出的SmaAt-UNet架构以及用于降水和云层现在预测的其他模型。第4节描述了进行的实验和获取的结果。第5节讨论了结果。最后，在第6节中，我们提出了一些结论性的意见。

2. 相关工作

一种基于深度学习的常见降水现在预报方法使用具有某种内存的神经网络，例如长短时记忆（LSTM）[18]。在标准的前馈模型中，输入数据以直接方式从一个时间步传递到下一个时间步。相比之下，LSTM是一种使输入信号保持在网络状态中多个时间步的网络，使网络能够记住过去的输入。这对于时间序列预测特别有用，因为过去的输入可能包含有关趋势的有用信息，进而有助于预测未来的值。

[6]的作者创建了一个ConvLSTM，用于在图像的时间序列任务中更好地捕捉时空相关性，优于其他方法。在此基础上，[7]的作者创建了一个时空LSTM(Predrnn)，增加了网络内部的内存连接数量，旨在实现空间信息的高效流动。该模型的内存功能和内存流程在另一种实现中进行了优化，添加了堆叠内存模块[19]。

降水现在预报的另一种方法已在[20]中描述。他们提出了一种基于众所周知的编码器-解码器架构UNet [16]的网络结构。与LSTM不同，UNet没有明确建模内存。它接受一个输入图像（或多个串联的图像）并输出单个分类地图。[20]的实现旨在对未来一小时内的四种不同降雨强度（<0:1mm/h，<1:0mm/h，<2:5mm/h，>2:5mm/h）进行分类。为此，将多个降水图（过去一小时的）串联在一起，用作UNet架构的输入。

在[9]的类似研究中，与[20]描述的模型不同，作者分类了512个类别，而不仅仅是四个，从而获得了更细的降雨强度分辨率。这与我们的方法类似；但是，与预测类别不同，我们的模型预测确切的降雨强度。

降水现在预报的常见基线是持续方法。持续模型使用序列的最后一个输入图像作为预测图像。这基于天气不会在从时间点t到t+1之间发生显着变化的假设。特别是在现在预报中，这个基线不容易被超越，因为图像之间的时间差非常短（例如2或5分钟），通常天气条件保持不变[1]。

最近，已经证明CNN中的注意力可以成为提高底层任务性能的非常有用的工具[21,22,23,24,25]。注意力是一种机制，它增强了所需信号并抑制了不需要的信号。这将网络引导到更关注对手头任务重要的特征。在我们提出的模型中，我们使用卷积块注意模块（CBAMs），这些模块按顺序对输入图像的通道和空间维度应用注意力[26]。CBAM的结果是一个加权特征图，它考虑了通道和输入图像的空间区域。

据我们所知，这是第一次在基于UNet的架构中包含CBAM机制。

在注意力的另一个应用中，[24]的作者为医学分割任务的UNet架构添加了注意力门。他们发现，与[16]的原始UNet实现相比，他们增强的模型获得了更好的结果。

网络中具有较少参数可以减少可能的过拟合风险，因为模型更简单，不能过于紧密地适应训练集的分布。这种简化的可能缺点是模型可能太简单，无法学习所需的任务。为了减少参数数量而不牺牲太多性能，许多最近的架构中使用了深度可分离卷积（DSC）[27,28,29,17,30]。DSC将常规卷积操作拆分为两个独立操作：深度卷积，然后是点卷积。这导致了较少的数学操作和较少的参数，与非分离卷积相比。[30]的作者创建了一个使用DSC而不是常规卷积的UNet，并且他们的模型的参数比原始UNet实现少了8倍。他们表明，他们的模型能够在医学分割任务上与UNet具有类似的性能[30]。

3. 方法论

3.1. 提出的 SmaAt-UNet 模型

我们在这里提出的模型是在 UNet 架构的基础上进行扩展和改进的[16]。如图1所示，UNet 架构包括编码器-解码器结构，呈现出 U 形状。编码器部分（对应于图1的左半部分）应用最大池化（红色箭头）和双卷积（蓝色箭头），分别将图像尺寸减半并将特征图数量加倍。编码器随后跟随相同数量的解码器（对应于图1的右半部分）。与UNet的原始实现一样，我们在这里也使用了四个编码器-解码器模块。

解码器由三个部分组成：双线性上采样操作（绿色箭头）以将特征图大小加倍，将生成的特征图与前一个编码器的输出通过跳过连接进行串联（灰色箭头），最后是双卷积以将特征图数量减半。跳过连接使模型能够使用输入的多个尺度来生成输出。最后，我们模型的最后一层是一个1x1卷积（紫色箭头），它输出一个表示网络预测值的单个特征图。

使用不同尺度的优势在于它们可以捕获输入中不同大小的对象，这对某些任务非常重要。通常，UNet被应用于分类或分割任务，在这些任务中，网络被训练以预测每个像素的类别。但是，我们将其应用于时间序列预测任务，在该任务中，网络必须为每个像素预测一个精确的值。

我们的新型 Small Attention-UNet（SmaAt-UNet）模型对原始UNet架构进行了两处修改。首先，我们建议将 CBAM 注意机制添加到编码器部分。其次，我们将常规卷积操作转换为深度可分离卷积。

如第2节所述，在CNN中使用注意力机制有助于网络集中关注输入的特定部分。对于我们的模型，我们使用卷积块注意模块来识别图像的跨通道和空间区域的重要特征[26]。在CBAM中，首先应用注意机制跨图像的通道，然后应用到空间维度。

CBAM被放置在第一个双卷积之后和每个编码器之后，以放大重要特征并抑制不重要的特征在相应的图像尺度上（图1中的黄色箭头）。重要的是，编码器的输入是来自前一个编码器的经过卷积和下采样的图像，而不是应用了注意机制的图像。这样，原始图像特征会被保留直到最后一个编码器。注意模块仅输入到它们通过跳过连接连接到的解码器的相应上采样部分。

沿着[27, 17]的思路，我们在模型中使用了深度可分离卷积，以减少参数数量。具体来说，我们用深度可分离卷积替换了原始UNet模型的所有卷积。但是，在卷积块注意模块中，我们仍然应用常规卷积。

3.2. 其他模型

为了进行比较，我们还训练了其他 UNet 架构，它们要么没有我们提出的两种修改，要么只有其中一种。因此，本研究共比较了四个模型，即原始 UNet、带有 CBAM 的 UNet、带有 DSC 的 UNet 和我们提出的模型。表1显示了这些模型的参数比较。当查看标准 UNet 架构和我们提出的修改后的 UNet 架构时，可以看到后者的参数明显更少，即约为 4 百万个参数，而标准 UNet 约为 17 百万个参数。在我们的 PyTorch 实现中，我们使用了每层两个内核的 DSC。

3.3. 训练

所有四个前述的模型都进行了最多 200 个 epochs 的训练。我们采用了早停止标准，当验证损失在最后的 15 个 epochs 中没有增加时，会停止训练过程。早停止标准在所有的训练迭代中都满足，因此最大的 200 个 epochs 从未达到。此外，我们使用了一个学习率调度程序，当验证损失在四个 epochs 内没有增加时，将学习率降低到前一个学习率的十分之一。初始学习率设置为 0.001，并使用默认值的 Adam 优化器 [31]。训练是在一台配备 8GB VRAM 的单个 NVidia 2070 Super 上进行的。

3.4. 模型评估

本研究中使用的损失函数是输出图像与地面实况图像之间的均方误差（MSE）。MSE 的计算如下：

除了均方误差（MSE）之外，我们还计算了不同的性能评估指标，例如精度、召回率（检测概率）、准确率、F1 分数、关键成功指数（CSI）、假警报率（FAR）和 Heidke 技能分数（HSS）。对于降水图数据集，这些分数是针对大于 0.5mm/h 阈值的降雨计算的。为此，我们使用该阈值将预测输出和目标图像的每个像素转换为布尔掩码。对于云覆盖数据集，数据已经进行了二值化。从中，可以计算真正例（TP）（预测=1，目标=1）、假正例（FP）（预测=1，目标=0）、真负例（TN）（预测=0，目标=0）和假负例（FN）（预测=0，目标=1）。随后，可以计算 CSI、FAR 和 HSS 指标，计算公式如下：

0.5mm/h 的阈值（第一个数据集）是根据 [6, 32] 的研究选择的，它用于区分雨水和无雨水。

4.实验

4.1. 降水图数据集

我们使用了荷兰皇家气象研究所（Koninklijk Nederlands Meteorologisch Instituut, KNMI）提供的降水数据作为第一个数据集来训练和比较我们的模型。该数据集包含了荷兰及邻近国家地区在过去四年（2016-2019年）中每5分钟间隔的降雨图。总共，该数据集包含约42万张降雨图。这些数据是由位于荷兰的De Bilt（52.10°N，5.18°E，44mMSL）和Den Helder（52.96°N，4.79°E，51mMSL）的两个C波段多普勒天气雷达站生成的。为了获取降雨图，这两台雷达会围绕垂直轴进行360度的方位扫描，仰角分别为0.3°、1.1°、2.0°和3.0°。雷达的其他参数可在表2中找到。此外，降雨图已经经过雨量计调整，更多细节可以在[33]中找到。我们将数据集分为训练集（2016-2018年）和测试集（2019年）。此外，在每次训练迭代中，我们通过随机选择训练集的10%来创建验证集。 原始的降雨图的尺寸为765x700，一个像素对应于一平方公里上过去五分钟内的累积降雨量。降雨量以百分之一毫米为单位表示为整数值。例如，数值12表示过去五分钟降雨量为0.12毫米。 作为数据准备步骤，我们将训练集和测试集的值除以训练集中最高出现的值以进行数据归一化。此外，我们裁剪了图像，只使用原始降水图的子集（如图2所示）。之所以这样做，是因为原始图像的许多像素都没有数据值，因为原始图像比雷达的最大范围要大（请参见图2左侧面板的白色边缘）。雷达范围内的区域呈圆形，直径为421个像素，相当于421公里。在以保留整个雷达图像的方式裁剪图像时，仍然存在许多没有数据值的像素（图2中间面板的白色角）。由于使用没有数据值的像素来训练神经网络更加困难，因此我们应用了额外的中心裁剪，裁剪掉了288个像素（图2右侧面板）。 模型的输入是一系列12张降水图，这些图像沿通道维度堆叠。这对应于过去一小时的天气观测（12 x 5分钟）。输出是比最后一个输入图像晚30分钟的降水图。因此，网络的任务是预测288 x 288降雨图中每个像素的精确降雨强度，预测未来30分钟的情况。 数据集包含许多降雨很少或没有降雨的降水图。因此，为了避免使网络偏向于预测零值，我们创建了两个附加数据集，其目标图像中具有最少的降雨像素。其中一个数据集中，目标图像中至少有20%的降雨像素，另一个数据集中至少有50%的降雨像素；我们将它们分别称为NL-20和NL-50。这两个数据集中的样本数量必然明显小于原始数据集，但它们更适用于模型的用例，即降雨预测。可以在表3中找到这三个数据集不同样本大小的比较。可以通过向作者提交请求来获取这些数据集。

我们在目标图像中至少有50%的降雨的数据集（NL-50）上训练了模型。这应该使受训网络的重点集中在降雨实例上。类似的做法也出现在[6]的研究中，他们选择了三年数据集中前97个多雨的日子进行训练。此外，这还可以用来作为模型的泛化性能的指标。更具体地说，我们可以将其用作模型的泛化性能指标。受过训练的模型从未见过此测试数据集中的任何降水图。此外，由于主要是多降雨的降水图，模型可能会有偏向于预测更多降雨的情况。因此，模型在此测试集上的性能可能不如其性能最接近的数据。

4.2. 云量数据集

作为第二个数据集，我们使用了在[34]中引入的云量数据集来比较提出的模型。该数据集包含了二值图像，每个像素表示云的位置，其中1表示像素上有云覆盖，0表示像素上没有云。所有图像的尺寸为256x256像素，包括法国的空间区域。此外，每个数据样本包含四张输入的二值图像和六张作为地面真实输出的二值图像。每个数据样本中的图像相隔15分钟，这导致了1小时的输入和1小时30分钟的输出，总共每个样本为2小时30分钟。有关该数据集的更多细节可以在[34]中找到。如图3所示，这个数据集的示例样本。

由于云覆盖数据集包含只有0和1的二值值图像，我们不需要应用任何数据归一化或图像裁剪。云覆盖数据集以其原始形式用于训练、评估和测试。网络的任务是预测每个像素上是否存在云的概率。

5.结果&讨论

5.1. 降雨图数据集上的评估

我们在降雨图数据集（NL-50）上获得的结果列在表4中。请注意，MSE是在将模型预测值恢复到原始降雨强度（mm/5min）后计算的。此外，我们还计算了MSE值除以两个不同数据集的平均像素值。结果是归一化MSE（NMSE），通过它我们可以在不同数据集之间进行公平的误差值比较。

获得的结果表明，在降雨图数据集上，每个模型都大幅超越了常见的持续性基线。这值得注意，因为如前所述，由于输入和目标之间的时间变化很小，很难在现在预测中超越这个基线。

我们发现，在降雨图数据集上，将我们提出的两种修改，即DSC和CBAM，添加到UNet架构中，与原始UNet实施相比，改变了模型的性能。一方面，单独实施每个修改都略微降低了性能。另一方面，然而，我们提出的模型SmaAt-Unet，它将这两个修改都融入到普通UNet中，比仅与每个修改之一结合的UNet表现更好。应该注意的是，仅配备CBAM的UNet在降雨图数据集上的MSE最高，为0.0171。至于我们的第二个修改，即用DSC替换常规卷积，结果则更加混合。一方面，带有DSC的UNet的性能比原始UNet模型差（分别为0.0127和0.0122）。但是，它仍然比带有CBAM的UNet模型性能更好。另一方面，重要的是要注意，将常规卷积替换为DSC将网络的模型大小减小到原始UNet的四分之一。

图4显示了降雨图数据集上降水现在预测模型输出的示例。与地面真实图像（左上角面板）相比，所有模型的预测降水图都相当模糊。造成这种情况的一个原因是使用MSE作为引导损失，它对模糊图像有偏向[35]。对模糊性的偏向是由于，鉴于基于输入序列的未来帧有许多可能性，模型正在尝试通过预测最接近所有可能结果的值来保持误差较低[36]。或者，正如Babaeizadeh等人所说的，“使用均方误差损失函数训练的模型会独立地生成每个像素所有可能性的期望值，这可能是固有的模糊”[37]。

此外，从图4可以看出，SmaAt-Unet能够更好地捕捉到强降雨群（左下角）的发展。UNet with DSCs预测的水平延伸范围过大。带CBAM的UNet做得更好，但预测的值过于保守。UNet产生了与SmaAt-UNet类似的输出，但无法很好地预测左侧降雨群的降水垂直扩散。

此外，我们计算了我们的模型在降雨图数据集上的性能多个指标。获得的分数也在表4中列出。这张表显示，尽管在大多数分数中，原始UNet实施在性能上表现最好，但我们的SmaAt-UNet在七项指标中的六项中表现次优。因此，尽管只有原始UNet四分之一的参数，SmaAt-Unet仍然能够接近UNet的性能。这有助于研究实验室和个人，他们没有太多的计算能力来进行这些计算密集型的计算。这反过来可以加速基于雷达的短期降雨预测的发展。

为了测试模型的泛化能力，我们使用了我们数据集的另一部分，如第4节中所描述的NL-20。这个测试集的MSE、NMSE和得分列在表5中。正如表中所示，结果与表4中的结果相似。具体来说，当对模型进行排名时，我们可以看到原始UNet实现在大多数指标中表现最佳，而我们的SmaAt-UNet在几乎所有指标中排名第二。这意味着尽管模型没有看到太多的降雨较少的输入，但UNet和SmaAt-UNet能够从训练时可用的有限数据中最好地外推。对于这个数据集的MSE较低的一个解释是，降水图的更多值接近零（因为降雨很少），因此如果模型也预测小值，则不会通过很大的边际增加整体MSE。因此，使用NMSE进行比较是更好的指标，因为它考虑了像素值分布。在这里，UNet略优于SmaAt-UNet，但两者的性能都比其他比较的模型好得多。

图5显示了在降雨图数据集上，来自编码器模块的注意力部分的示例特征图。该图表明，网络的注意力图学会关注输入序列的特定部分，展示了注意机制的学习效果。行代表编码器的不同阶段，可以通过每行中分辨率的减小来看出。此外，可以看出注意力特征图关注输入的不同特征。例如，在第一行中，一些特征图关注左下角的降雨簇（图2和8），而其他特征图则关注几乎没有降雨的部分（图4、5和7）。底部一行显示了SmaAt-UNet的最后一个编码器阶段的特征图，其分辨率为18 x 18。正如底部一行的图像所示，这种低分辨率使网络能够识别粗糙的模式，比如图像底部的雨簇（图2、3、5和7）。

5.2. 云覆盖数据集上的评估

我们还在云覆盖数据集上训练了讨论过的四个UNet-based模型。多个性能指标的结果列在表6中。云覆盖数据集包含具有二进制值的样本，因此在这里我们不计算NMSE错误。此外，数据集格式设置为每个样本四个输入图像和六个输出图像。

从表6中可以看出，UNet与CBAM具有最低的MSE得分。然而，所有模型的结果都是可比较的。在对所有四个模型在云覆盖数据集上获得的MSE的最高值和最低值之间的差异只有0.0019的情况下，这个小差异尤其值得一提，对于SmaAt-UNet和Unet with DSC来说，它们的参数约为原始Unet或带CBAM的Unet的1/4。

在表6中的大多数报告指标中，Unet with CBAM获得了最佳得分。在两种情况下，Unet with CBAM与另一个模型相媲美，即F1分数与Unet相媲美，FAR分数与Unet with DSC相媲美。对于这个数据集，即模型任务是预测每个像素上的云概率在0和1之间，所提出的Unet与卷积块注意力模块的结合在大多数情况下达到了最佳分数。然而，所提出的SmaAt-UNet在具有原始Unet和Unet with CBAM的大约1/4的参数的情况下，获得了非常相似的性能。

请注意，这些翻译是根据您提供的文本进行的，某些领域的专业术语可能需要根据上下文和目标读者进行微调，以确保准确性和可理解性。如果您需要更多细节或有任何其他疑问，请随时提出。

需要注意的是，在云覆盖数据集的情况下，与降雨图数据集相比，评分之间的差异较小，因为云覆盖数据集的数据值是二进制的。对于云覆盖数据，评估模型预测的值在0和1之间（云的存在或不存在），这导致与对降雨图数据集的评估相比差异较小（请参阅表4和表5）。

6.总结

在本论文中，我们提出了SmaAt-UNet，它是UNet架构的一个更小、更具关注性的版本。已经证明它在降雨现在预测任务上的性能与比它大得多的类似架构相当。开发小型和高效的神经网络，如SmaAt-UNet，使它们能够在智能手机上应用。例如，创建一个具有多个经过训练的SmaAt-Unet的应用程序，每个具有不同的预测时间，可以根据用户的请求使用最新的可用数据进行降水预测。此外，创建能源高效的架构，如SmaAt-UNet，可以减少碳足迹。对培训神经网络所需资源的考虑是迈向可持续机器学习实践的关键步骤。