基于多尺度混合注意力LSTM雷达回波外推方法

刘瑞华,高 翔,邹洋杨

(重庆理工大学 两江人工智能学院,重庆 401135)

0 引言

短临预报,即短时效(0～2 h)内的天气变化预报,在交通、农业、军事、畜牧业和旅游业等都起到至关重要的作用,对于城市防汛预警更是无可替代。雷达回波序列图像因其自身具有较高的时间和空间相关性的特点,常作为临近预报的有效工具,而如何利用历史雷达回波图像进行临近预报特别是对强对流天气的预测是一个热点问题[1-2]。

近年来,短临预报采用的单体质心法主要有雷暴识别跟踪分析的临近预报方法[3]以及风暴单体识别与跟踪算法[4]。随着研究的深入,虽然单体质心法经过不断完善[5-8],也取得了一些进展,但是计算量庞大,以及泛化能力差的缺点使其无法在临近预报上取得更好的效果。交叉相关外推算法最早由Rinehart等[9]提出,曹伟华等[10]先对连续雷达图像进行相关分析,找到相邻时次图像的位移矢量,然后预测用于外推雷达回波移动的位置和强度。陈明轩等[11]指出,基于回波移动矢量需要获得平滑和真实的回波移动矢量才能进行比较准确的外推,该算法存在外推时效较短的缺陷。符式红等[12]构造了集合的线性交叉相关方法对多普勒雷达回波进行外推,能够反映降水系统所对应回波的移动方向和速度,考虑到多时刻对流单体的演变过程,随着外推时间步长的增加,回波强度与实况相比慢慢偏弱,无法进行较为精准的预测。光流法最早应用于计算机视觉领域,通过利用图像序列中在时间域上的变化以及相邻帧间的相关性来找到上一帧和当前帧之间的对应关系。曹春燕等[13]提到对准静止型局地加强型降水,光流法预报效果较低,该算法存在一定的局限性,且光流法进行外推预报的可用性随着时间的延长而降低。

随着深度学习的愈发流行,并结合雷达回波资料时空相关的数据特点,学者们将循环神经网络(reccurent neural network,RNN)算法应用到临近预报,并取得了一些研究进展[14-18]。Shi等[19]通过把卷积神经网络(convolution neural network,CNN)中的卷积层引入长短时记忆网络(long short time memory network,LSTM),提出了卷积长短时记忆网络(convolution LSTM network,Conv-LSTM),主要是在LSTM运算单元中加入卷积,让模型既能对时序信息建模,同时又兼备卷积的局部感知和减少参数量的优点。Shi等[20]考虑到雷达回波具有时空易变的特性,提出了对连续多帧的质点运动矢量进行追踪,捕捉其在时间序列上的空间变化的轨迹GRU(trajectory GRU,Traj-GRU)模型。该模型对自然运动条件下的观测物体仍能保持较高的预测精度,并提出了加权损失函数作为指标判断模型对不同阈值下的雷达回波预测能力。Wang等[21]采用改进的时空长短时记忆网络,在网络中增加了一个并行的专用于保留图片空间特征的通道,进而提出了预测RNN (predictive recurrent neural network,Pred-RNN)模型。该模型不会在经过多层的LSTM运算后丢失大量空间信息,同时也解决时空长短时记忆网络层级式结构中存在的梯度消失问题。Wang等[22]提出了重记忆模型,将LSTM中的遗忘门变成2个级联的LSTM结构,使网络具有捕捉时间序列中非平稳特征的能力。Wu等[23]提出了运动RNN的通用视频预测模型,该模型通过在时空上把运动变化分解为整体变化与瞬时变化,并利用运动GRU(gated recurrent unit,GRU)单元统一建模,另外该模型可平衡序列预测中的移动与非移动部分信息。当然,卷积神经网络应用于视频预测,降水预报也愈发流行[24-25]。Han等[26]利用Unet进行连续降水预报,使用更少的参数能够得到和LSTM相差无几的效果。Trebing等[27]提出在Unet上增加注意力模块的小型注意力Unet模型也取得了不俗的预测效果。Zhang等[28]针对中高强度的雷达回波预测问题,把视觉Transformer[29]和LSTM结合,提出了局部注意力预测模型,并很好地解决了此问题。然而,上述模型也存在几个小缺陷：① 由于LSTM递归的特点,使其无法像CNN一样搭建不同图像尺寸的层级式网络,仅仅依靠LSTM中的几个门,即便结合了卷积神经网络,在空间特征提取上仍存在着不足。② Transformer中图片划分块的大小也影响着Transformer本身的准确度,如细粒度的块能够使Transformer性能更好,而粗粒度的块则会使Transformer在自注意力计算时能得到更多的全局信息,另一方面LSTM无法做到同时提取细粒度和粗粒度的块信息。

为了解决上述问题,受Pred-RNN模型与注意力机制的影响,本文提出了多尺度混合注意力网络模型(multi-scale mixed attention network,MMA-Net)。该模型主要由多个层级的多尺度混合注意力MMA构成,与Pred-RNN不同之处在于该单元中具有不同尺度的回波图像输入,并通过一种全新的混合注意力特征提取方法(cross attention block,CAB)着重加强了LSTM对空间特征的提取,实现了在transformer中对细粒度和粗粒度块的有效结合。

1 MMA-Net

雷达回波外推是关于时间和空间序列的预测问题,基于一系列历史雷达回波图片去预测未来降水。假设在由M行N列组成的一张M×N的空间区域内观测降水,其中每一个单位网格内存在P个随时间变化的观测值,对应在t时刻下的一张雷达回波图像可用张量Xt∈R,P×M×N进行表示,此问题可以被描述为：

(1)

式中：R代表观测的特征域,以所观测的t时刻为当前时刻,向前取连续的J帧图像作为输入图像序列,预测未来长度为K的雷达回波图像。

1.1 Pred-RNN

图1 卷积循环神经网络

1.2 CAB模块

2017年,Vaswani 等[29]首次提出了Transformer模型,目前已广泛应用于计算机视觉任务,如图像识别、图像分类和分割等,在架构上可以与CNN结合起来混合使用,从而达到更好的效果。但Transformer与其他模型结合时,其分块大小也一定程度上影响着本身的准确度,如细粒度的块能够使Transformer性能更好,而粗粒度的块则会使自注意力计算时能得到更多的全局信息。

由于现实因素,通常情况下,中高回波强度在整张回波图像中比例极小,如图2所示,分块的大小会影响模型对不同回波强度的预测精确度,对于中低等回波强度,小的分块(如图2(c))往往能够提供连续信息,而大的分块(如图2(b))则会使模型对中高回波强度更敏感。

基于此,本文设计了CAB模块。该模块首先对同一时刻不同尺寸的输入图像块进行特征提取,然后对粗粒度和细粒度的图像子块进行融合,力争解决LSTM模型无法对不同粒度图片信息进行提取的缺陷,以提高模型对强对流天气的预测准确度,如图3所示。

图2 不同的分块大小

图3 CAB模块

为了解决LSTM在空间特征提取方面的不足,本模型在输入端上就当前时刻输入与隐藏层输入分别添加了辅助分支,CAB则是针对2个不同尺度输入的图像特征提取模块,该模块由2次特征融合构成。在图3(a)中分为分支(branch,B)以及混合注意力(cross attention,CA)2个部分,在B部分,较小尺寸的辅助分支给模型带来细粒度信息,双路分支结合主要是提高模型对中高回波强度信息的捕捉能力,同时通过上采样使辅助分支与主分支共同提取特征,使模型关注中低回波强度的信息。在CA部分,粗粒度信息和细粒度信息的注意力混合更着重于全局信息提取。考虑到不同块大小仅单次采样Transformer中的权重矩阵会造成不平衡的影响,把图3(a)的输出作为第二次特征融合的辅助分支,并与主分支进行再次融合构建不同的权重矩阵,如图3(b)所示,完成了对不同尺度特征的提取。

1.3 MMA

图4 MMA-LSTM

MMA-LSTM的数据流主要由4个级联组成,具体计算过程为：

(2)

(3)

(4)

(5)

(6)

式中：CAB计算流程见图3,计算中所使用的激活函数皆为双曲正切函数,同时,与Pred-RNN模型不同的是本模型中在t时刻输入以及隐藏层输入上分别增加了一个辅助分支,通过上述CAB模块提高模型的空间特征提取能力,并多出一个用于增加历史回波信息提取的通道自注意力(self attention,SA)。SA模块即视觉Transformer中的自注意力模块,通过SA可以计算出历史回波图像的一个权重矩阵,从而得出过去的哪张回波图像能对下一时刻的预测提供更大帮助,进而提升模型对回波外推的长时间建模的能力。

1.4 MMA-Net

本文采用基于ST-LSTM改进的MMA模块组成的神经网络模型作为网络架构,具体结构如图5所示。

图5 MMA-Net

模型延续了ST-LSTM的之字形结构设计,该设计能够保留前一帧的信息同时应用于后一帧的预测。由于回波图像外推的重点是对时间和空间的序列预测,所以与其他模型不同的是本模型在横向与纵向的数据流上面增加了部分的历史回波输入X,通过历史回波输入,能够更好地抓住序列之间的时间关系进行建模。而在每一个MMA模块中,通过添加辅助分支并利用CAB融合方法,有效地把Transformer中的大尺度块和小尺度块的优势进行结合,同时本模块通过辅助分支与原输入的多尺度输入,采用全新的特征融合方法构建了一种类似CNN的层级式特征提取。

2 仿真实验

数据来源于深圳气象局与阿里巴巴联合承办的CIKM Analyti-Cup 2017比赛,是深圳及周边地区的多普勒雷达观测数据,作为公开数据集,该数据集中每一个样本包含15张雷达回波序列图片,间隔为6 min,共包括90 min的观测数据。该数据集共包括10 000个雷达回波序列作为训练集,4 000个雷达回波序列作为测试集,本实验在训练集中随机选取2 000个序列作为验证集。每一张图片空间分辨率为1 km×1 km,每一个像素值范围0～255,通过Z-R关系可以把像素值转为对应的雷达回波强度。Z-R关系如下：

(7)

本实验利用5张雷达回波图像预测未来连续10张图像,即预测未来1 h的降水情况。本文实验设备处理器型号为AMD 5900X,24 GB内存,NVIDIA RTX 3090 GPU,使用Pytorch实现本文模型。网络模型训练时,每一个块大小为64,学习率为0.000 1,尺寸大小设置为4,采用Adam优化器进行优化。实验中,损失函数为MSE和MAE的加权组合,即,

Loss=MSE+0.2×MAE

(8)

2.1 评价指标

本次实验对测试集检验采用分阈值的评估方法,使用5、20、40 dBz作为判别阈值,评估指标主要包括CSI和HSS,具体计算公式如下。

(9)

(10)

式中：TP、FP、TN、FN分别表示真阳性、假阳性、真阴性、假阴性。

2.2 预测评估实验

为了评估MMA-Net模型性能效果,测试了CSI和HSS评估指标在3个不同阈值(5、20、40 dBz)情形下,与其他先进模型进行比较,如表1所示。通过比较发现,除了20 dBz下的CSI指标,所提出的MMA-Net模型在CSI和HSS各阈值情形下都取得了最高的指标值,表明模型取得了较好的效果。

表1 测试集评估指标

雷达回波预测问题是时空预测问题,在本实验中,为更好地体现随时间关系不同模型的预测效果,绘制了不同阈值下的HSS和CSI曲线,如图6所示。

通过图6中的HSS和CSI曲线可以看出,本模型在5 dBz和20 dBz下对未来10 帧即1 h的回波预测效果超过其他模型,而在40 dBz的前6帧,即36 min前本模型要优于其他模型,而在第6帧即36 min后RAP-Net要优于本模型,分析是雷达回波图像的预测比较依赖于历史信息,因为RAP-Net模型能够保留所有的历史输入,同时其对每张图像提取特征不够充分,则导致RAP-Net在中低回波强度以及高回波强度平均成绩上略低。

图6 不同回波强度下的CSI和HSS曲线

图7、图8为测试样本的可视化图像,其中第一行图片表示真实回波图像,最下面颜色表格代表不同的雷达回波强度,回波强度越高则表示该区域发生极端天气情况的几率越高。由于本次实验是利用5张回波图像对未来连续10张回波图像进行预测,给出了不同模型的可视化预测结果。可以看出,Conv-LSTM和Conv-GRU模型对中高回波强度的预测误差很大。Traj-GRU从第7帧开始,不能很好地预测35 dBz以上的回波强度,Pred-RNN以及Rap-Net在6帧之前对40 dBz以上的回波强度预测效果不佳,6帧之后则无法很好地对中等回波强度20 dBz进行预测。恰如上述曲线图中,6帧之后对于高回波强度40 dBz的预测本模型稍逊于Rap-Net,因为指标的评定依赖于完整10张回波图像预测,所以在总体上即包含低、中、高3种回波强度预测下本模型优于其他模型。

图7 第340序列测试样本的可视化图像(第一行：真实回波图像)

图8 第933序列测试样本的可视化图像(第一行：真实回波图像)

2.3 消融实验

为了讨论CAB模块在MMA-Net模型中的作用,本小节将对此进行消融实验。为了便于简单了解CAB模块在不同层而造成的模型差异,在此对模型做一些说明：MMA-Net-F,MMA-Net-I和MMA-Net-H模型分别表示仅在第一层加入辅助分支和CAB模块,在输入中加入辅助分支和CAB模块,以及在隐藏状态加入辅助分支和CAB模块,如表2、表3所示。

表2 消融实验(学习率为10-4)

表3 消融实验(学习率为10-5)

通过表2和表3可以看出,在MMA-Net模型中,如果仅在第一层增加CAB模块,此模型可用于高强度的雷达回波外推预测。而如果仅在隐藏层增设CAB模块,此模块会更适合于中低强度的雷达回波外推预测。通过消融实验,MMA-Net模型中的CAB模块对于提升不同程度雷达回波外推预测能力有较好的效果,可以根据需求适当调整CAB模块位置的设计。

3 结论

雷达回波外推的精确预测对城市防洪防汛预警具有重要意义。针对CIKM Analyti-Cup 2017雷达回波数据集,提出了MMA-Net模型,在LSTM长期序列建模以及更好的空间特征表示问题上取得了较好的实验仿真效果,同时也取得了较高的评估指标值及较好的可视化效果。解决模型对硬件资源需求较高、参数量过大和效率不高的问题,将是后续研究内容。