临近降水预报是指某一地区在较短的时间内 (一般来说是0–6 h) 的降雨预报. 准确的临近降水预报不仅可以为人们的日常生活提供便利, 还可以减少降雨造成的损失. 目前降水预报常用的业务系统是数值天气预报和光流法^[1,2], 但由于数值天气预报基于流体力学和热力学方程, 计算复杂度高, 其精度严重依赖于初始条件. 而光流法需要假设雷达回波的亮度是恒定的. 实际上, 雷达回波的亮度会因降水强度变化、衰减等各种因素而发生变化. 因此在预测降水事件的起始和衰减方面仍然表现不佳. 为了提高精度, 深度学习已成为降水临近预报的重要发展方向. 基于深度学习的临近降水预报一般通过雷达回波外推实现. 首先利用过去雷达回波数据预测未来雷达回波数据, 然后利用$ Z {\textit{-}} R $关系 ($Z = 10\log a + 10b\log R$) 将预测雷达回波数据转化为降水强度数据, 实现降水强度临近预报, 其中$R$表示降雨量mm/h, $Z$表示雷达回波反射率dBZ, $a$和$b$是基于特定区域统计数据计算的两个常数.

目前, 深度学习临近降水预报网络中主要使用卷积神经网络 (convolutional neural network, CNN) 和循环神经网络 (recurrent neural network, RNN). 这些方法以端到端方式训练神经网络模型, 通过学习大量的复合雷达观测数据的统计规律进行预测^[3]. 然而, 雷达回波的演变涉及整体的平移变化和局部的空间形变^[4], 如图1所示, 橙色和红色区域明显的向右下角方向移动, 同时局部区域也在时刻发生着变化, 如, 黑色框内红色雷达回波区域范围变得更大且强度增强. 以前的方法尝试在一个统一的框架中对这两种变化进行建模, 但是由于大气环境的多样性, 雷达回波的变形、积累或消散非常复杂, 一个统一的框架难以同时处理时间和空间上的复杂变化.

基于RNN的方法, 如ConvLSTM等虽然在LSTM中集成了卷积操作来建模时空变化, 但是在反向传播过程中容易出现梯度值逐渐减小并趋近于零或梯度值不断增大造成梯度消失或梯度爆炸的问题. 除此之外, 这些方法通常是一阶马尔可夫模型^[5], 对于长序列的处理能力相对较弱, 容易出现短期记忆问题, 即过早丢弃先前的信息, 难以捕捉到长期的上下文信息. 而基于CNN的方法在处理序列数据时对于时间关系的建模能力有一定的局限性. 传统的卷积操作主要用于处理空间信息, 而在处理时间序列数据时, 卷积操作无法直接捕捉到时间上的依赖和顺序关系.

为了更好地处理复杂的雷达回波演化过程, 本文提出了一种基于Halo注意力机制^[6]的双阶段临近降水预报网络, 该网络采用Halo注意力机制作为基本构建单元构建了可学习光流模块和特征重建模块, 旨在分别对雷达回波的运动趋势和空间形变进行建模, 提高深度学习方法的时空建模能力, 提供准确且具有细粒度特征的预测结果. 该网络将预测雷达回波演变分为两个阶段. 第1阶段利用可学习光流模块学习雷达回波序列中降水云体的运动模式, 提供对降水的位置和移动趋势的粗略预测. 第2阶段利用特征重建模块对可学习光流模块的粗略预测结果进行特征细化重建. 总体而言, 本文提出的方法为临近降水方向提供了一种新的研究思路, 具体的创新点总结如下.

(1) 将降水的演变过程分为两个阶段, 并分别设计了可学习光流模块和特征重建模块, 这种分阶段的设计可以充分利用时间和空间信息, 提高降水预测的准确性和稳定性, 并为未来的研究提供了新的思路和方法.

(2) 使用Halo注意力机制作为所提出网络的基本构建单元, 使得网络能够自适应地分配注意力权重, 更好地捕捉到雷达回波序列中的时间和空间相关性.

1 相关工作 1.1 光流法

光流法是计算机视觉领域中常用的方法, 用于估计图像序列中像素的运动, 其基本原理是通过比较相邻帧之间的像素亮度差异来计算光流场. 光流场表示了每个像素在两帧之间的位移情况, 即其运动方向和速度. 通过对整个图像应用光流法, 可以得到每个像素的光流向量场, 从而揭示了整个图像中物体的运动情况. 在气象领域, 光流法^[7–10]逐渐被应用于雷达回波的跟踪, 并取得了良好的实际效果. 例如, 香港天文台开发了多尺度光流变分法, 它结合了光流法和变分法的思想, 并在不同尺度上进行计算, 用于雷暴单体预报和定量降水预报. 此外, 香港天文台还进一步提出了实时雷达回波光流变分法, 该方法采用了空间平滑法并且允许调整光流变分参数^[9]. 另外, Ayzel等人^[10]开发了一种名为rainymotion的临近降水模型, 该模型基于多种光流算法计算降水的运动信息, 进而进行降水的预测.

1.2 深度学习方法

2015年, Shi等人^[11]将临近降水预报作为一个时空序列预测问题, 并提出了ConvLSTM网络, 该网络用卷积结构取代了FC-LSTM中的全连接结构进行预测, 这为临近降水预报开辟了新的方向. Wang等人^[12,13]提出了一种称为PredRNN的时空预测方法, PredRNN额外增加了时空记忆单元, 并通过“之”字形结构将它们连接起来, 使特征在空间和时间上传播来帮助建模时空信息. Wang等人^[14]提出了一种名为MIM的模型, 该模型创新性地利用相邻循环状态之间的差分信号来模拟时空动力学的稳态和非稳态特征, 并在多个时空序列数据集上取得了最佳的预测结果. Niu等人^[15]提出了一种双阶段的时空细化网络2S-STRef, 该模型第1个阶段生成粗粒度的预测结果, 第2个阶段对第1个阶段生成的预测结果进行细化. Ravuri等人^[3]提出了一种深度生成模型, 以消除较长预测时间内临近预报结果中的模糊问题. Wu等人^[16]基于物理世界运动可以分解为瞬态变化和运动趋势的原理提出了MotionRNN框架来捕捉运动中的复杂变化, 同时通过MotionGRU单元对瞬态变化和运动趋势进行建模.

最近一些研究^[17–21]将CNN扩展到降水临近预报中, 其中具有代表性的CNN模型是UNet及其变体. Zhuang等人^[17]提出了ST-CNN来充分利用时空信息, 从大量的降水数据中自动学习降水的基本模式, 进行极端降水预测. Ayzel等人^[18]提出了RainNet, 其采用UNet架构从大量的雷达数据中学习时空演变来提供精确的临近降水预测. Trebing等人^[19]通过为UNet配备卷积块注意力模块和深度可分离卷积提出了SmaAt-UNet. Yang等人^[20]结合了Transformer和UNet模型提出了AA-TransUNet, 并通过结合卷积块注意力模块和深度可分离卷积在降低模型参数的同时提升模型性能.

2 双阶段临近降水预报网络 2.1 网络架构

本文采用端到端的学习方法来生成预测雷达回波序列. 如图2所示, 所提出的网络架构由可学习光流模块和特征重建模块组成, 这两个模块都是神经网络, 可以实现非线性建模. 为了可以更好地捕捉到降水事件的空间相关性和时序特征, 这两个模块采用Halo注意力机制作为基本的构建单元构成. 此外, 为了更好地提取和保留图像特征, 两个网络都采用了“U”型结构并通过跳跃连接来拼接低级和高级雷达回波特征.

每帧雷达回波图像可以表示为${X_t} \in {R^{H \times W \times C}}$, 其中$H$、$W$和$C$分别表示高度、宽度和通道数. 第1个阶段, 可学习光流模块对5帧连续的雷达回波图像$ {X_{t - 4:t}} $进行时间相关性进行建模生成预测光流场${\hat V_{t + 1:t + 10}}$并通过warp操作生成10帧预测的粗粒度雷达回波图像${\tilde X_{t + 1:t + 10}}$. 第2个阶段, 特征重建模块将从历史序列$ {X_{t - 4:t}} $学习到的空间变化模式对粗粒度预测结果${\tilde X_{t + 1:t + 10}}$进行细化重建, 生成具有细粒度特征的预测雷达回波序列${\hat X_{t + 1:t + 10}}$. 这种在一个统一的网络中通过分别建模运动趋势和空间变化的方法, 可以充分地建模时空变化情况.

2.2 Halo注意力机制

随着注意力机制的发展, 在计算机视觉领域, 注意力已是一种常用的算法, 并且已被应用于不同的任务中. 注意力机制的公式可以表示为式 (1) 所示:

$ {\textit{Attention(}}Q{\text{, }}K{\text{, }}V{\text{) = {\textit{Softmax}}}}\left(\frac{{Q{K^{\mathrm{T}}}}}{{\sqrt {{d_k}} }}\right)V $

(1)

其中, $Q,\; K, \; V \in {R^{H \times W \times C}}$分别表示Query、Key和Value, $k$为注意力头的数量, ${d_k} = C/k$是注意力头的维度. 虽然注意力机制基于内容的交互和更强的捕获图像中远程依赖关系的能力使得重要信息获得了更多的关注. 但是, 注意力机制通常需要大量计算, 导致计算复杂性增加. 如果采用全局的注意力, 对于较大的图像就会导致计算量非常大. 如果采用局部的注意力, 由于每次滑动窗口的结果在显存中得不到释放, 可能就会导致内存不足的问题. Halo注意力将整张图片分为多个块, 并对每个块进行局部注意力, 从而解决了上述问题.

除此之外, 一个地点的降水受距离它较远地点的降水的影响很小. 因此, 要预测给定地点的降水情况, 只需要关注附近区域的降水情况, 可以忽略那些太远地区的降水情况. 而Halo注意力恰恰遵循了局部处理的原则, 不仅可以提升降雨预测的效果, 而且解决了在高分辨率下注意力机制的计算成本过大的问题. 如图3所示, Halo注意力机制主要分为3个步骤: 首先将整个图片分为多个块, 其次, 如果每次操作时只考虑每个块包含的信息, 则不可避免地会导致信息丢失. 因此, 在计算每个块之前, 在每个块之外, 利用原图像的信息进行填充 (简称Halo), 使得每个块的感受野可以适当放大, 可以关注到更多的信息. 最后, 对每个块分别执行局部自注意力.

另外, 传统的注意力模型在进行下采样的时候通常使用平均池化或最大池化. Halo注意力机制的下采样层可以对每个块分别进行采样, 然后对这些采样后的信息进行Halo注意力操作, 从而达到下采样的效果. 因为这里是对每个块的信息进行采样后计算, 而不是对每个块内的信息进行计算, 所以这样的操作是不会影响模型精度的, 同时还将下采样层中的FLOPs降低了4倍.

2.3 可学习光流模块

雷达回波图序列的演化是一个复杂的过程, 由于复杂的大气环境, 雷达回波不仅会发生平移运动, 雷达回波的形状也可能会迅速膨胀、收缩或变化, 这与其他时空序列预测任务有极大的不同. 因此, 本文将降水预测任务分解为两个阶段, 并设计了可学习的光流模块和特征重建模块来处理这两个阶段的不同任务. 可学习光流模块旨在通过学习雷达回波序列中降水云体的运动模式, 粗略推测未来一段时间内降水云体的位置情况^[22]. 与以往的基于CNNs的光流生成网络相比, 本文提出的可学习光流模块引入了Halo注意力机制, 这使得网络能够更好地关注重要的特征或动态地调整特征的权重, 以适应不同的运动场景. 这种动态调整特征权重的能力对于处理复杂的大气环境下的雷达回波图序列非常关键, 因为雷达回波可能会发生各种形状和运动模式的变化. 此外, 通过使用Halo注意力机制, 光流模块能够自适应地学习到不同区域和运动模式之间的相关性, 并将更多的注意力放在具有更显著运动的区域上.

如图4所示, 可学习光流模块由上侧编码器和下侧解码器构成. 编码器和解码器均有5层, 蓝色的Halo注意力机制不使用下采样层, 其输出的特征图尺寸不会改变, 而绿色的Halo注意力机制使用下采样层. 在本文中, 使用下采样层后, 其输出的特征图尺寸会缩小为原尺寸的1/4. 对于编码器, 第1层将连续的5帧雷达回波图像$ {X_{t - 4:t}} $的通道维度扩充至64通道. 其余4层通过Halo注意力机制的下采样层减小空间维度来逐步提取更抽象的特征. 除编码器第5层不进行通道维度扩充外, 中间3层每一层都将通道维度翻倍.

对于解码器, 除第1层不进行通道维度压缩, 其余各层均将通道维度压缩至输入维度的一半. 编码器和解码器通过跳跃连接将特征在通道维度上进行拼接, 如图4中红色箭头所示. 为了更好地建模运动趋势并改善光流估计的性能, 编码器和解码器之间拼接的是不同尺度的特征信息, 由于拼接的特征图尺寸不一致, 因此在拼接前需要对尺寸较小的特征图进行边缘填充. 此外, 解码器最后一层使用1$ \times $1卷积将特征维度压缩至20, 从而最终生成连续的10帧未来雷达回波图像的光流场${\hat V_{t + 1:t + 10}}$. 所生成的每个光流场具有两个通道, 表示了图像中每个像素点的运动方向和速度. 可学习光流模块每一阶段的具体特征尺寸如表1所示. 最后, 通过式 (2) 对输入雷达回波序列的最后一帧$ {X_t} $进行warp操作得到下一时刻的雷达回波图$ {\tilde X_{t + 1}} $, warp操作可以将雷达图像根据光流场信息进行变换, 重复此过程直至生成预测的未来10帧连续的雷达回波图像${\tilde X_{t + 1:t + 10}}$.

$ {\tilde X_{t + 1}} = {\textit{warp}}({X_t}, {\hat V_{t + 1}}) $

(2)

2.4 特征重建模块

特征重建模块利用深度学习强大的表征学习能力学习雷达回波复杂的空间变化模式, 并基于可学习光流模块提供的粗粒度特征信息进行空间特征重建. 如图5所示, 本文提出的特征重建模块由上侧编码器和下侧解码器构成. 编码器的输入为在通道维度上拼接的5帧连续的历史雷达回波图像$ {X_{t - 4:t}} $和可学习光流模块生成的10帧粗粒度预测结果${\tilde X_{t + 1:t + 10}}$. 历史雷达回波图像中包含了雷达回波的过去状态和空间外观变化. 而可学习光流模块的粗粒度预测结果提供了雷达回波在时间上的运动趋势. 特征重建模块编码器通过结合运动趋势信息可以更好地理解和建模雷达回波的空间形态变化. 此外, 编码器通过逐渐抽象的过程使得编码器能够捕捉到不同尺度上的特征, 从而更好地描述雷达回波的演变过程, 包括膨胀、收缩和形状变化等. 解码器的主要任务是将编码器学习到的空间特征变化模式应用到多尺度的粗粒度预测序列上, 并通过非线性上采样逐步恢复特征的空间维度. 整个过程采用了循序渐进的方式, 通过逐渐对不同尺度的粗粒度特征进行空间细节重建, 并将这些不同尺度的信息融合在一起, 可以得到具有细粒度特征的预测序列${\hat X_{t + 1:t + 10}}$.

与可学习光流模块结构图中的颜色意义相同, 蓝色的Halo注意力机制不使用下采样层, 其输出的特征图尺寸不会更改, 而绿色的Halo注意力机制使用下采样层. 编码器第1层保持图像尺寸不变, 通道维度扩充至64, 随后3层均会将输入序列进行下采样, 图像尺寸缩减为缩小为原尺寸的1/4, 通道维度增加一倍. 最后一层保持图像尺寸和通道维度不变. 在非线性上采样前, 可学习光流模块的预测结果经过自适应池化为解码器不同阶段特征的尺寸, 这种多尺度的预测结果在通道维度与编码器的特征拼接, 可以提供多尺度的雷达回波特征信息用于特征重建. 最后一层编码器的输出特征图经1$ \times $1卷积操作进行维度压缩后得到最终的预测雷达回波序列. 特征重建模块每一阶段的具体特征尺寸如表2所示. 总体而言, 特征重建模块通过学习历史雷达回波序列的空间变化模式并对多尺度的雷达回波特征信息进行特征重建来生成具有细粒度细节的预测雷达回波图.

3 实验分析 3.1 数据集

本文利用信息与知识管理国际会议 (CIKM AnalytiCup 2017) 挑战赛数据集来评估所提出的模型. 它包含14000个序列, 其中8000个序列作为训练集, 2 000个序列作为验证集, 4000个序列作为测试集. 训练集和验证集中的数据是连续两年的, 而测试集中的数据是在下一年的^[23]. 每个序列包含15帧雷达回波图像, 间隔6 min, 仰角分别为0.5 km、1.5 km、2.5 km和3.5 km. 在本文中, 选择3.5 km仰角的雷达回波图像来训练和测试所提出的模型. 在每个序列中, 前5帧雷达回波图被作为模型输入, 后10帧作为预测输出. 每张图像分辨率为101$ \times $101. 每个像素代表1 km$ \times $1 km的面积, 所有像素的取值范围为 [0–255].

3.2 评估标准

作为预处理步骤, 式 (3) 将每个像素的像素值转换为雷达反射率:

$ Z=\frac{像素值}{255}\times 90-10 $

(3)

在评估指标方面, 采用阈值法将预测回波图与观察值进行转换. 如果该值大于给定阈值, 则将相应值设置为1; 否则设置为0. 本文计算了真阳性预测TP (预测=1, 真值=1)、假阳性预测FP (预测=1, 真值=0)、真阴性预测TN (预测=0, 真值=0) 和假阴性预测FN (预测=0, 真值=1) 的数量. 此外, 本文使用了5 dBZ、20 dBZ和40 dBZ这3个阈值, 并用式 (4) 和式 (5) 计算了关键成功指数 (critical success index, CSI) 和海德克技能得分 (Heidke skill score, HSS) 来评估结果:

$ {\textit{CSI}} = \frac{{TP}}{{TP + FN + FP}} $

(4)

$ {\textit{HSS}} = \frac{{2(TP \times TN - FN \times FP)}}{{(TP + FN)(FN + TN) + (TP + FP)(FP + TN)}} $

(5)

这两个指标广泛应用于降水预报的准确性测评中. CSI指标关注模型预测的降水事件与真实观测到的降水事件之间的一致性, 即成功预测的降水事件的比例, 其取值范围从0到1, 值越接近1表示预测结果越准确. HSS评估的是排除随机预测后的正确预测比例, 它衡量模型的预测能力相对于随机猜测的改进程度, 其取值范围从−1到1, 值越接近1表示预报结果越准确. 除此之外, 结构相似性 (structural similarity index, SSIM) 和均方误差 (mean squared error, MSE) 也作为本文的评估指标. SSIM考虑了图像的亮度、对比度和结构等方面的信息, 是基于人眼感知的图像质量评价方法之一, 可以用来评估图像的质量和清晰度. MSE用于衡量两幅图像之间差异的指标. 它计算了两幅图像对应像素之间的平均平方差, 从而反映了它们在像素级别上的差异程度.

3.3 实验参数设置

所有实验均在 PyTorch 框架下实现, 并且在NVIDIA GeForce RTX 3090上运行. 所提出的网络使用 Adam 优化器在训练过程中进行优化, 学习率设置为0.000 3, 批次大小设置为16. 此外, 损失函数使用的是平均绝对误差 (mean absolute error, MAE). 为防止模型过拟合, 该网络在训练过程中采用了早停机制. 原始输入${X_t}$维度为 (101, 101, 1), 为方便数据操作, 本文使用0填充图像四周, 将输入数据维度改为(128, 128, 1) 后输入到模型中进行预测. 预测结果去掉四周填充的部分后, 变为初始输入维度 (101, 101, 1).

3.4 实验结果

为了验证所提出模型的优越性, 本文与6种主流模型进行了比较, 包括ConvLSTM^[11]、ConvGRU^[24]、PredRNN^[12]、PredRNN++^[25]、MIM^[14]和PredRNN_V2 ^[13]. 如表3所示, 最佳结果以粗体标出, 所提出的模型在所有指标上都表现出了优势. 为了更直观地展示本文网络的提升效果, 本文绘制了HSS和CSI指标的折线图, 展示了不同阈值下各模型随预测时间增加的表现. 如图6所示, 所提出的网络在所有阈值下均明显领先于先前的模型, 这意味着其具有更高的准确性, 进一步凸显了其卓越性能. 此外, 图7还展示了各模型的SSIM折线图, 结果显示所提出的网络获得了最高的得分, 这说明其预测结果更加逼真, 与观测到的雷达回波图像更为接近. 为了验证所提出网络的稳定性, 本文还基于生成的4000个雷达回波序列计算了各项指标的方差. 为了更加直观地观察结果, 我们将MSE指标的方差缩小了10⁷倍. 如表4所示, 可以明显看出本文所提出的方法在各项指标的方差方面表现最佳, 这充分展示了所提出模型的稳定性.

为了进一步展示所提出模型的实际预测效果, 本文随机选择了一个雷达回波序列, 并绘制了所提出模型以及其他主流模型的预测效果图, 如图8所示. 通过对比可以明显看出, 与先前的模型相比, 所提出的模型在预测雷达回波的位置和强度方面表现更为准确. 特别是在黄色区域, 先前的模型错误地预测了其位置, 甚至对雷达回波的强度进行了高估. 而所提出的网络则更加接近观测结果, 其预测结果在雷达回波的形状和强度方面更加精确. 综上所述, 通过对比各项指标的折线图、方差结果以及预测效果图的展示, 可以充分证明所提出模型相对于其他主流模型在准确性和可靠性方面的优越性, 为相关领域的研究和应用提供了有力的支持.

3.5 消融实验

为了评估所使用的Halo注意力机制在网络中的有效性以及本文所提出的双阶段预测网络的实际提升效果, 本文进行了消融实验, 分别对Halo注意力机制、单阶段模型进行了独立验证. 首先, 本文移除可学习光流模块, 用3$ \times $3卷积替代特征重建模块中的Halo注意力机制, 并以此作为基准 (Baseline) 来验证Halo注意力机制的效果. 其次, 本文使用应用了Halo注意力机制的特征重建模块作为单阶段预测模型. 根据表5的结果显示, 仅使用特征重建模块的Baseline网络在预测效果上表现最差, 这表明Halo注意力机制对于提高网络性能具有重要作用. 单阶段网络的效果明显提升, 进一步验证了Halo注意力机制的有效性. 值得注意的是, 当引入可学习光流模块, 构建了双阶段网络时, 预测效果进一步提升. 这说明可学习光流模块能够在时序上建模雷达回波的运动趋势, 其融合到预测过程中增强了模型对降水演变的理解和预测能力. 综上所述, 消融实验的结果验证了Halo注意力机制的有效性, 并证明了所提出的双阶段临近降水预报网络能够更好地利用空间和时间信息, 提高了对临近降水的预测能力, 为临近降水预报研究提供了有力的支持.

4 结论与展望

本文提出了一种基于Halo注意力机制的双阶段临近降水预报网络, 该网络由可学习光流模块和特征重建模块组成. 这两个模块是由Halo注意力机制作为基本构建模块构建的, 并将降水预测的过程分为两个阶段. 在第1个阶段, 可学习光流模块通过估计雷达回波序列中的运动趋势生成粗粒度的预测结果. 在第2个阶段, 特征重建模块对历史雷达回波序列的空间外观进行建模并基于学习到的空间外观变化模式对粗粒度的预测结果进行特征重建, 从而实现对未来降水的准确预测. 这种双阶段的预测方法能够更好地捕捉雷达回波的时空演变特征, 提高降水预报的准确性和细节展示. 实验结果表明, 所提出的网络在应对复杂的雷达回波演变时具有更高的准确性和清晰度.