基于多尺度混合注意力LSTM雷達(dá)回波外推方法

劉瑞華,高 翔,鄒洋楊

(重慶理工大學(xué) 兩江人工智能學(xué)院,重慶 401135)

0 引言

短臨預(yù)報,即短時效(0～2 h)內(nèi)的天氣變化預(yù)報,在交通、農(nóng)業(yè)、軍事、畜牧業(yè)和旅游業(yè)等都起到至關(guān)重要的作用,對于城市防汛預(yù)警更是無可替代。雷達(dá)回波序列圖像因其自身具有較高的時間和空間相關(guān)性的特點,常作為臨近預(yù)報的有效工具,而如何利用歷史雷達(dá)回波圖像進(jìn)行臨近預(yù)報特別是對強(qiáng)對流天氣的預(yù)測是一個熱點問題[1-2]。

近年來,短臨預(yù)報采用的單體質(zhì)心法主要有雷暴識別跟蹤分析的臨近預(yù)報方法[3]以及風(fēng)暴單體識別與跟蹤算法[4]。隨著研究的深入,雖然單體質(zhì)心法經(jīng)過不斷完善[5-8],也取得了一些進(jìn)展,但是計算量龐大,以及泛化能力差的缺點使其無法在臨近預(yù)報上取得更好的效果。交叉相關(guān)外推算法最早由Rinehart等[9]提出,曹偉華等[10]先對連續(xù)雷達(dá)圖像進(jìn)行相關(guān)分析,找到相鄰時次圖像的位移矢量,然后預(yù)測用于外推雷達(dá)回波移動的位置和強(qiáng)度。陳明軒等[11]指出,基于回波移動矢量需要獲得平滑和真實的回波移動矢量才能進(jìn)行比較準(zhǔn)確的外推,該算法存在外推時效較短的缺陷。符式紅等[12]構(gòu)造了集合的線性交叉相關(guān)方法對多普勒雷達(dá)回波進(jìn)行外推,能夠反映降水系統(tǒng)所對應(yīng)回波的移動方向和速度,考慮到多時刻對流單體的演變過程,隨著外推時間步長的增加,回波強(qiáng)度與實況相比慢慢偏弱,無法進(jìn)行較為精準(zhǔn)的預(yù)測。光流法最早應(yīng)用于計算機(jī)視覺領(lǐng)域,通過利用圖像序列中在時間域上的變化以及相鄰幀間的相關(guān)性來找到上一幀和當(dāng)前幀之間的對應(yīng)關(guān)系。曹春燕等[13]提到對準(zhǔn)靜止型局地加強(qiáng)型降水,光流法預(yù)報效果較低,該算法存在一定的局限性,且光流法進(jìn)行外推預(yù)報的可用性隨著時間的延長而降低。

隨著深度學(xué)習(xí)的愈發(fā)流行,并結(jié)合雷達(dá)回波資料時空相關(guān)的數(shù)據(jù)特點,學(xué)者們將循環(huán)神經(jīng)網(wǎng)絡(luò)(reccurent neural network,RNN)算法應(yīng)用到臨近預(yù)報,并取得了一些研究進(jìn)展[14-18]。Shi等[19]通過把卷積神經(jīng)網(wǎng)絡(luò)(convolution neural network,CNN)中的卷積層引入長短時記憶網(wǎng)絡(luò)(long short time memory network,LSTM),提出了卷積長短時記憶網(wǎng)絡(luò)(convolution LSTM network,Conv-LSTM),主要是在LSTM運算單元中加入卷積,讓模型既能對時序信息建模,同時又兼?zhèn)渚矸e的局部感知和減少參數(shù)量的優(yōu)點。Shi等[20]考慮到雷達(dá)回波具有時空易變的特性,提出了對連續(xù)多幀的質(zhì)點運動矢量進(jìn)行追蹤,捕捉其在時間序列上的空間變化的軌跡GRU(trajectory GRU,Traj-GRU)模型。該模型對自然運動條件下的觀測物體仍能保持較高的預(yù)測精度,并提出了加權(quán)損失函數(shù)作為指標(biāo)判斷模型對不同閾值下的雷達(dá)回波預(yù)測能力。Wang等[21]采用改進(jìn)的時空長短時記憶網(wǎng)絡(luò),在網(wǎng)絡(luò)中增加了一個并行的專用于保留圖片空間特征的通道,進(jìn)而提出了預(yù)測RNN (predictive recurrent neural network,Pred-RNN)模型。該模型不會在經(jīng)過多層的LSTM運算后丟失大量空間信息,同時也解決時空長短時記憶網(wǎng)絡(luò)層級式結(jié)構(gòu)中存在的梯度消失問題。Wang等[22]提出了重記憶模型,將LSTM中的遺忘門變成2個級聯(lián)的LSTM結(jié)構(gòu),使網(wǎng)絡(luò)具有捕捉時間序列中非平穩(wěn)特征的能力。Wu等[23]提出了運動RNN的通用視頻預(yù)測模型,該模型通過在時空上把運動變化分解為整體變化與瞬時變化,并利用運動GRU(gated recurrent unit,GRU)單元統(tǒng)一建模,另外該模型可平衡序列預(yù)測中的移動與非移動部分信息。當(dāng)然,卷積神經(jīng)網(wǎng)絡(luò)應(yīng)用于視頻預(yù)測,降水預(yù)報也愈發(fā)流行[24-25]。Han等[26]利用Unet進(jìn)行連續(xù)降水預(yù)報,使用更少的參數(shù)能夠得到和LSTM相差無幾的效果。Trebing等[27]提出在Unet上增加注意力模塊的小型注意力Unet模型也取得了不俗的預(yù)測效果。Zhang等[28]針對中高強(qiáng)度的雷達(dá)回波預(yù)測問題,把視覺Transformer[29]和LSTM結(jié)合,提出了局部注意力預(yù)測模型,并很好地解決了此問題。然而,上述模型也存在幾個小缺陷：① 由于LSTM遞歸的特點,使其無法像CNN一樣搭建不同圖像尺寸的層級式網(wǎng)絡(luò),僅僅依靠LSTM中的幾個門,即便結(jié)合了卷積神經(jīng)網(wǎng)絡(luò),在空間特征提取上仍存在著不足。② Transformer中圖片劃分塊的大小也影響著Transformer本身的準(zhǔn)確度,如細(xì)粒度的塊能夠使Transformer性能更好,而粗粒度的塊則會使Transformer在自注意力計算時能得到更多的全局信息,另一方面LSTM無法做到同時提取細(xì)粒度和粗粒度的塊信息。

為了解決上述問題,受Pred-RNN模型與注意力機(jī)制的影響,本文提出了多尺度混合注意力網(wǎng)絡(luò)模型(multi-scale mixed attention network,MMA-Net)。該模型主要由多個層級的多尺度混合注意力MMA構(gòu)成,與Pred-RNN不同之處在于該單元中具有不同尺度的回波圖像輸入,并通過一種全新的混合注意力特征提取方法(cross attention block,CAB)著重加強(qiáng)了LSTM對空間特征的提取,實現(xiàn)了在transformer中對細(xì)粒度和粗粒度塊的有效結(jié)合。

1 MMA-Net

雷達(dá)回波外推是關(guān)于時間和空間序列的預(yù)測問題,基于一系列歷史雷達(dá)回波圖片去預(yù)測未來降水。假設(shè)在由M行N列組成的一張M×N的空間區(qū)域內(nèi)觀測降水,其中每一個單位網(wǎng)格內(nèi)存在P個隨時間變化的觀測值,對應(yīng)在t時刻下的一張雷達(dá)回波圖像可用張量Xt∈R,P×M×N進(jìn)行表示,此問題可以被描述為：

(1)

式中：R代表觀測的特征域,以所觀測的t時刻為當(dāng)前時刻,向前取連續(xù)的J幀圖像作為輸入圖像序列,預(yù)測未來長度為K的雷達(dá)回波圖像。

1.1 Pred-RNN

圖1 卷積循環(huán)神經(jīng)網(wǎng)絡(luò)

1.2 CAB模塊

2017年,Vaswani 等[29]首次提出了Transformer模型,目前已廣泛應(yīng)用于計算機(jī)視覺任務(wù),如圖像識別、圖像分類和分割等,在架構(gòu)上可以與CNN結(jié)合起來混合使用,從而達(dá)到更好的效果。但Transformer與其他模型結(jié)合時,其分塊大小也一定程度上影響著本身的準(zhǔn)確度,如細(xì)粒度的塊能夠使Transformer性能更好,而粗粒度的塊則會使自注意力計算時能得到更多的全局信息。

由于現(xiàn)實因素,通常情況下,中高回波強(qiáng)度在整張回波圖像中比例極小,如圖2所示,分塊的大小會影響模型對不同回波強(qiáng)度的預(yù)測精確度,對于中低等回波強(qiáng)度,小的分塊(如圖2(c))往往能夠提供連續(xù)信息,而大的分塊(如圖2(b))則會使模型對中高回波強(qiáng)度更敏感。

基于此,本文設(shè)計了CAB模塊。該模塊首先對同一時刻不同尺寸的輸入圖像塊進(jìn)行特征提取,然后對粗粒度和細(xì)粒度的圖像子塊進(jìn)行融合,力爭解決LSTM模型無法對不同粒度圖片信息進(jìn)行提取的缺陷,以提高模型對強(qiáng)對流天氣的預(yù)測準(zhǔn)確度,如圖3所示。

圖2 不同的分塊大小

圖3 CAB模塊

為了解決LSTM在空間特征提取方面的不足,本模型在輸入端上就當(dāng)前時刻輸入與隱藏層輸入分別添加了輔助分支,CAB則是針對2個不同尺度輸入的圖像特征提取模塊,該模塊由2次特征融合構(gòu)成。在圖3(a)中分為分支(branch,B)以及混合注意力(cross attention,CA)2個部分,在B部分,較小尺寸的輔助分支給模型帶來細(xì)粒度信息,雙路分支結(jié)合主要是提高模型對中高回波強(qiáng)度信息的捕捉能力,同時通過上采樣使輔助分支與主分支共同提取特征,使模型關(guān)注中低回波強(qiáng)度的信息。在CA部分,粗粒度信息和細(xì)粒度信息的注意力混合更著重于全局信息提取?？紤]到不同塊大小僅單次采樣Transformer中的權(quán)重矩陣會造成不平衡的影響,把圖3(a)的輸出作為第二次特征融合的輔助分支,并與主分支進(jìn)行再次融合構(gòu)建不同的權(quán)重矩陣,如圖3(b)所示,完成了對不同尺度特征的提取。

1.3 MMA

圖4 MMA-LSTM

MMA-LSTM的數(shù)據(jù)流主要由4個級聯(lián)組成,具體計算過程為：

(2)

(3)

(4)

(5)

(6)

式中：CAB計算流程見圖3,計算中所使用的激活函數(shù)皆為雙曲正切函數(shù),同時,與Pred-RNN模型不同的是本模型中在t時刻輸入以及隱藏層輸入上分別增加了一個輔助分支,通過上述CAB模塊提高模型的空間特征提取能力,并多出一個用于增加歷史回波信息提取的通道自注意力(self attention,SA)。SA模塊即視覺Transformer中的自注意力模塊,通過SA可以計算出歷史回波圖像的一個權(quán)重矩陣,從而得出過去的哪張回波圖像能對下一時刻的預(yù)測提供更大幫助,進(jìn)而提升模型對回波外推的長時間建模的能力。

1.4 MMA-Net

本文采用基于ST-LSTM改進(jìn)的MMA模塊組成的神經(jīng)網(wǎng)絡(luò)模型作為網(wǎng)絡(luò)架構(gòu),具體結(jié)構(gòu)如圖5所示。

圖5 MMA-Net

模型延續(xù)了ST-LSTM的之字形結(jié)構(gòu)設(shè)計,該設(shè)計能夠保留前一幀的信息同時應(yīng)用于后一幀的預(yù)測。由于回波圖像外推的重點是對時間和空間的序列預(yù)測,所以與其他模型不同的是本模型在橫向與縱向的數(shù)據(jù)流上面增加了部分的歷史回波輸入X,通過歷史回波輸入,能夠更好地抓住序列之間的時間關(guān)系進(jìn)行建模。而在每一個MMA模塊中,通過添加輔助分支并利用CAB融合方法,有效地把Transformer中的大尺度塊和小尺度塊的優(yōu)勢進(jìn)行結(jié)合,同時本模塊通過輔助分支與原輸入的多尺度輸入,采用全新的特征融合方法構(gòu)建了一種類似CNN的層級式特征提取。

2 仿真實驗

數(shù)據(jù)來源于深圳氣象局與阿里巴巴聯(lián)合承辦的CIKM Analyti-Cup 2017比賽,是深圳及周邊地區(qū)的多普勒雷達(dá)觀測數(shù)據(jù),作為公開數(shù)據(jù)集,該數(shù)據(jù)集中每一個樣本包含15張雷達(dá)回波序列圖片,間隔為6 min,共包括90 min的觀測數(shù)據(jù)。該數(shù)據(jù)集共包括10 000個雷達(dá)回波序列作為訓(xùn)練集,4 000個雷達(dá)回波序列作為測試集,本實驗在訓(xùn)練集中隨機(jī)選取2 000個序列作為驗證集。每一張圖片空間分辨率為1 km×1 km,每一個像素值范圍0～255,通過Z-R關(guān)系可以把像素值轉(zhuǎn)為對應(yīng)的雷達(dá)回波強(qiáng)度。Z-R關(guān)系如下：

(7)

本實驗利用5張雷達(dá)回波圖像預(yù)測未來連續(xù)10張圖像,即預(yù)測未來1 h的降水情況。本文實驗設(shè)備處理器型號為AMD 5900X,24 GB內(nèi)存,NVIDIA RTX 3090 GPU,使用Pytorch實現(xiàn)本文模型。網(wǎng)絡(luò)模型訓(xùn)練時,每一個塊大小為64,學(xué)習(xí)率為0.000 1,尺寸大小設(shè)置為4,采用Adam優(yōu)化器進(jìn)行優(yōu)化。實驗中,損失函數(shù)為MSE和MAE的加權(quán)組合,即,

Loss=MSE+0.2×MAE

(8)

2.1 評價指標(biāo)

本次實驗對測試集檢驗采用分閾值的評估方法,使用5、20、40 dBz作為判別閾值,評估指標(biāo)主要包括CSI和HSS,具體計算公式如下。

(9)

(10)

式中：TP、FP、TN、FN分別表示真陽性、假陽性、真陰性、假陰性。

2.2 預(yù)測評估實驗

為了評估MMA-Net模型性能效果,測試了CSI和HSS評估指標(biāo)在3個不同閾值(5、20、40 dBz)情形下,與其他先進(jìn)模型進(jìn)行比較,如表1所示。通過比較發(fā)現(xiàn),除了20 dBz下的CSI指標(biāo),所提出的MMA-Net模型在CSI和HSS各閾值情形下都取得了最高的指標(biāo)值,表明模型取得了較好的效果。

表1 測試集評估指標(biāo)

雷達(dá)回波預(yù)測問題是時空預(yù)測問題,在本實驗中,為更好地體現(xiàn)隨時間關(guān)系不同模型的預(yù)測效果,繪制了不同閾值下的HSS和CSI曲線,如圖6所示。

通過圖6中的HSS和CSI曲線可以看出,本模型在5 dBz和20 dBz下對未來10 幀即1 h的回波預(yù)測效果超過其他模型,而在40 dBz的前6幀,即36 min前本模型要優(yōu)于其他模型,而在第6幀即36 min后RAP-Net要優(yōu)于本模型,分析是雷達(dá)回波圖像的預(yù)測比較依賴于歷史信息,因為RAP-Net模型能夠保留所有的歷史輸入,同時其對每張圖像提取特征不夠充分,則導(dǎo)致RAP-Net在中低回波強(qiáng)度以及高回波強(qiáng)度平均成績上略低。

圖6 不同回波強(qiáng)度下的CSI和HSS曲線

圖7、圖8為測試樣本的可視化圖像,其中第一行圖片表示真實回波圖像,最下面顏色表格代表不同的雷達(dá)回波強(qiáng)度,回波強(qiáng)度越高則表示該區(qū)域發(fā)生極端天氣情況的幾率越高。由于本次實驗是利用5張回波圖像對未來連續(xù)10張回波圖像進(jìn)行預(yù)測,給出了不同模型的可視化預(yù)測結(jié)果。可以看出,Conv-LSTM和Conv-GRU模型對中高回波強(qiáng)度的預(yù)測誤差很大。Traj-GRU從第7幀開始,不能很好地預(yù)測35 dBz以上的回波強(qiáng)度,Pred-RNN以及Rap-Net在6幀之前對40 dBz以上的回波強(qiáng)度預(yù)測效果不佳,6幀之后則無法很好地對中等回波強(qiáng)度20 dBz進(jìn)行預(yù)測。恰如上述曲線圖中,6幀之后對于高回波強(qiáng)度40 dBz的預(yù)測本模型稍遜于Rap-Net,因為指標(biāo)的評定依賴于完整10張回波圖像預(yù)測,所以在總體上即包含低、中、高3種回波強(qiáng)度預(yù)測下本模型優(yōu)于其他模型。

圖7 第340序列測試樣本的可視化圖像(第一行：真實回波圖像)

圖8 第933序列測試樣本的可視化圖像(第一行：真實回波圖像)

2.3 消融實驗

為了討論CAB模塊在MMA-Net模型中的作用,本小節(jié)將對此進(jìn)行消融實驗。為了便于簡單了解CAB模塊在不同層而造成的模型差異,在此對模型做一些說明：MMA-Net-F,MMA-Net-I和MMA-Net-H模型分別表示僅在第一層加入輔助分支和CAB模塊,在輸入中加入輔助分支和CAB模塊,以及在隱藏狀態(tài)加入輔助分支和CAB模塊,如表2、表3所示。

表2 消融實驗(學(xué)習(xí)率為10-4)

表3 消融實驗(學(xué)習(xí)率為10-5)

通過表2和表3可以看出,在MMA-Net模型中,如果僅在第一層增加CAB模塊,此模型可用于高強(qiáng)度的雷達(dá)回波外推預(yù)測。而如果僅在隱藏層增設(shè)CAB模塊,此模塊會更適合于中低強(qiáng)度的雷達(dá)回波外推預(yù)測。通過消融實驗,MMA-Net模型中的CAB模塊對于提升不同程度雷達(dá)回波外推預(yù)測能力有較好的效果,可以根據(jù)需求適當(dāng)調(diào)整CAB模塊位置的設(shè)計。

3 結(jié)論

雷達(dá)回波外推的精確預(yù)測對城市防洪防汛預(yù)警具有重要意義。針對CIKM Analyti-Cup 2017雷達(dá)回波數(shù)據(jù)集,提出了MMA-Net模型,在LSTM長期序列建模以及更好的空間特征表示問題上取得了較好的實驗仿真效果,同時也取得了較高的評估指標(biāo)值及較好的可視化效果。解決模型對硬件資源需求較高、參數(shù)量過大和效率不高的問題,將是后續(xù)研究內(nèi)容。