基于多尺度混合注意力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長期序列建模以及更好的空間特征表示問題上取得了較好的實驗仿真效果,同時也取得了較高的評估指標值及較好的可視化效果。解決模型對硬件資源需求較高、參數量過大和效率不高的問題,將是后續研究內容。