一種深度學習系統(tǒng)可以從高精度的氣候模擬中識別出極端天氣模式

近年來極端天氣所造成的影響越來越大，不久前的山竹和飛燕帶來的損失讓我們對自然的力量仍然記憶猶新。如何更好的識別極端天氣，做好防災(zāi)減災(zāi)工作是氣象人員面臨的巨大挑戰(zhàn)。

而最近，為了幫助識別極端天氣的潛在模式，預(yù)測出有可能損害人民群眾生命財產(chǎn)安全的天氣模式，來自伯克利大學勞倫斯實驗室、橡樹嶺實驗室和英偉達的研究人員們，利用AI和超級計算機給出了答案。他們研發(fā)了一種深度學習系統(tǒng)可以從高精度的氣候模擬中識別出極端天氣模式，這種算法有望在未來幫助公眾盡早預(yù)知天氣，在天災(zāi)面前從容應(yīng)對。

研究人員們利用Tiramisu和DeepLabv3神經(jīng)網(wǎng)絡(luò)為極端天氣構(gòu)建了像素級的掩膜。在最近放出的論文里，研究人員詳細描述了在軟件架構(gòu)上的改進，包括輸入流程和訓練算法有效的提升了在超算上的大規(guī)模深度學習計算能力。Tiramisu網(wǎng)絡(luò)被大規(guī)模的部署在5300個P100的英偉達GPU上，實現(xiàn)了21PF/s的運算能力以及79%的并行效率。同時DeepLabv3則部署到了27360個V100GPU上，實現(xiàn)了325.8PF/s的速度，同時在單精度下實現(xiàn)了90.7%的并行效率。同時通過高效的使用FP16張量核心，半精度版本的DeepLabv3+網(wǎng)絡(luò)可以實現(xiàn)1.13EF/s的峰值運算和999.0 PF/s的持續(xù)運算能力。這些計算在目前排名第一的超算Summit上進行。

研究人員表示這一研究實現(xiàn)了很多記錄，不僅是在氣象科學領(lǐng)域第一次將深度學習用于解決如此大規(guī)模的氣候分割問題，同時在深度學習領(lǐng)域也是第一次將計算能力拓展到了Exa（艾10^18）量級。1.13EF/s的計算能力是深度學習首次跨越Exa的大門。值得一提的是，這一突破也獲得了今年的Gordon Bell獎，用于鼓勵研究團隊在超算領(lǐng)域的杰出成就。

下面讓我們來看看研究人員使用的網(wǎng)絡(luò)架構(gòu)，用于對天氣狀況進行像素級的分割。

這項工作使用改進DeepLabv3+網(wǎng)絡(luò)，用于對氣候數(shù)據(jù)進行分割。其中編碼器部分使用了ResNet-50的核心部分來提取輸入數(shù)據(jù)的高維特征，并使用空洞空間金字塔池化來處理較大的輸入分辨率。而其解碼器則被完整分辨率的統(tǒng)一操作取代，以便生成精確的邊界信息。在上圖中，深藍色框代表了標準的卷積操作，淺藍色代表了反卷積操作，而空洞卷積操作則由綠色代表。

研究人員們使用了CommunityAtmosphere Model (CAM5) 來生成100年的模擬數(shù)據(jù)，并用啟發(fā)式的算法來生成掩膜標簽。整個數(shù)據(jù)集包含了63k的高分辨率樣本，其中包含了水蒸氣、風力、降水、溫度和氣壓等，分割區(qū)域包含了熱帶氣旋、大氣長河和背景三種區(qū)域。

為了在超算上大規(guī)模的部署應(yīng)用，研究人員們測試了不同數(shù)量GPU的計算速度：

結(jié)果顯示目前的算法和虛線表示的結(jié)果十分接近。

下圖是最終對于全球范圍內(nèi)的預(yù)測掩膜結(jié)果，我們可以看到預(yù)測出了結(jié)果和黑色的標記結(jié)果十分接近。其中藍色為大氣長河紅色為熱帶氣旋。

這一工作在先前氣象模式分割的啟發(fā)下成功地將Tiramisu和DeepLabV3+網(wǎng)絡(luò)應(yīng)用到了高分辨率多變量的氣候數(shù)據(jù)中，并在對損失函數(shù)、優(yōu)化機制和網(wǎng)絡(luò)結(jié)構(gòu)的改進下，在定量和定性方面都實現(xiàn)了優(yōu)異的結(jié)果。同時，為了大規(guī)模部署算法研究人員們建立了系統(tǒng)級的優(yōu)化方式，包括了先進的數(shù)據(jù)存儲策略、優(yōu)化吞吐和層級通訊，使得峰值節(jié)點4560情況下利用27360個VoltaGPU實現(xiàn)了高效的并行計算和超高的算力（峰值1.13EF/s，均值999.0PF/s），這一研究拓展了開源工具Tensorflow和Horovod的使用，研發(fā)的系統(tǒng)已經(jīng)部署在了Summit超算中。這將對深度學習領(lǐng)域帶來諸多貢獻，同時也將促進超算平臺的發(fā)展。