当前位置: 首页 > 科技观察

速度提升45000倍,Nvidia利用傅立叶模型实现前所未有的天气预报精度

时间:2023-03-20 15:25:19 科技观察

现代数值天气预报(numericalweatherprediction,NWP)可以追溯到1920年代。如今,数值天气预报无处不在,有助于交通、物流、农业和能源生产等关键部门的经济规划。准确的天气预报通过提前通知极端事件挽救了无数生命。在过去的几十年里,天气预报的质量一直在稳步提高。最早的单点动力模拟数值天气预报是由英国科学家刘易斯·弗莱·理查森于1922年使用计算尺和对数表计算得出的,计算一次6小时(6小时)的大气预报需要六周时间。到1950年代,早期的电子计算机极大地提高了预报速度,使得运行预报的计算速度足以帮助预测未来的天气。除了更好的计算能力外,还可以通过更深入地了解小规模过程的物理特性和更高质量的大气观测,对小规模过程进行更好的参数化,从而改进天气预报。数据驱动的深度学习模型比SOTANWP模型便宜几个数量级,因此研究人员现在对开发此类方法来预测天气越来越感兴趣。许多研究试图通过建立数据驱动的模型和对气候模型输出、大气环流模型(GCM)、再分析产品或气候模型输出和再分析产品的组合进行培训来预测大气中的大规模环流。通过克服NWP模型中存在的模型偏差以及通过以低计算成本生成用于概率预测和数据同化的大型集合,数据驱动模型具有改善天气预报的巨大潜力。大多数数据驱动的天气模型都是使用低分辨率数据进行训练的,Rasp和Thuerey[2021b]中的分辨率通常为5.625°,Weyn等人中的分辨率通常为2°。[2020]。这些先前的尝试在预测一些粗粒度、低分辨率的大气变量方面取得了很好的效果。然而,粗化过程会导致关键的细粒度物理信息丢失。要使数据驱动模型真正产生影响,它们必须以与当前SOTA数值天气模型相同或更高的分辨率(大约0.1°)生成预报。在最近的一篇论文中,来自英伟达、劳伦斯伯克利国家实验室、密歇根大学安娜堡分校、莱斯大学等机构的研究人员开发了一种基于傅立叶的神经网络预测模型FourCastNet,该模型可以生成关键天气变量的全球数据驱动预测分辨率为1,000?,相当于赤道附近的空间分辨率约为30×30km,全球网格大小为720×1440像素。这使我们能够首次与欧洲中期天气预报中心(ECMWF)的高分辨率综合预报系统(IFS)模型进行直接比较。论文地址:https://arxiv.org/pdf/2202.11214.pdfFourCastNet在节点小时基础上比传统NWP模型快约45,000倍。FourCastNet在高分辨率下的这种数量级加速和前所未有的准确性使其能够以一小部分成本生成非常大的整体预测。FourCastNet极大地改进了概率天气预报,使用它可以在几秒钟内生成飓风、大气河流和极端降水等事件的大规模集合预报,从而实现更及时、更明智的灾难响应。此外,FourCastNet可靠、快速且廉价的近地表风速预报可以改善陆上和海上风电场的风能资源规划。训练FourCastNet所需的能量大约等于使用IFS模型(50个成员)生成10天预测所需的能量。然而,经过训练后,FourCastNet生成预测所需的能量比IFS模型少12,000倍。研究人员希望FourCastNet只训练一次,后续微调的能量消耗可以忽略不计。在实现技术方面,FourCastNet使用基于傅立叶变换的令牌混合方法[Guibasetal.,2022]和ViT主干[Dosovitskiyetal.,2021]。这种方法基于最近的傅立叶神经算子,这些算子以分辨率不变的方式学习,并成功地模拟了具有挑战性的偏微分方程,例如流体动力学。此外,他们选择ViT骨干网是因为它能够很好地模拟远程依赖关系。ViT和基于傅立叶的令牌方法的混合产生了SOTA高分辨率模型,可以解析细粒度特征并根据分辨率和数据集大小很好地扩展。据研究人员称,这种方法能够以真正前所未有的高分辨率训练高保真数据驱动模型。训练模型欧洲中期天气预报中心(ECMWF)提供了一个公开可用的合成数据集ERA5,本研究使用该数据集来训练FourCastNet。除了关注两个大气变量,即(1)地球表面以上10米的风速和(2)6小时总降水量外,该研究还预测了其他几个变量,包括几个位势高度、温度、风速和相对湿度,一些近地表变量,如表面压力和平均海平面压力等。为了为FourCastNet模型生成高分辨率预测,该研究使用自适应傅立叶神经算子(AFNO)模型。这种神经网络架构专为高分辨率输入而设计,并将深度学习(DL)的最新关键进展整合到单个模型中。也就是说,FourCastNet结合了傅里叶神经算子(FNO),它已被证明在具有挑战性的PDE系统建模中表现良好,并且还具有强大的ViT主干。目前,研究人员提出了几种降低计算复杂度的ViT变体,然而,AFNO模型的独特之处在于它将混合运算构建为连续的全局卷积,由FFT在傅里叶域中有效地实现,这允许灵活和可扩展的建模跨空间和通道维度的依赖性。通过这样的设计,空间混合复杂度降低到O(NlogN),其中N是图像块或标记的数量。这种缩放使AFNO模型非常适合0.25°分辨率的高分辨率数据,以及未来更高分辨率的研究。模型的计算过程:首先将720×1440lat-lon网格上的输入变量投影到patch的二维网格(h×w)上(patchsize为p×p,例如p=8),每个patch表示为一个d维的token。然后,补丁序列连同位置编码被馈送到一系列AFNO层。训练研究重点是预测地表风速和降水量,但复杂的大气系统包含多个变量之间的强非线性相互作用,如温度、地表气压、湿度、地表到平流层的水分含量等。为了模拟这些相互作用,该研究选择了变量(表1)来表示大气的瞬时状态。该研究将ERA5数据集分为三组,即训练、验证和测试数据集。训练数据集包含1979年至2015年的数据,验证数据集包含2016年和2017年的数据,测试数据集包含2018年及以后的数据。本研究使用张量X(kΔt)来表示变量,其中k是时间索引,Δt是训练数据集中连续快照之间的时间间隔。研究以ERA5数据集为真值,用X_true(kΔt)表示真值变量,Δt固定为6小时。训练包括预训练和微调两个阶段:在预训练阶段,使用训练数据集以有监督的方式训练AFNO模型,使其学习到从X(k)到X(k)的映射X(k+1);在微调阶段,研究从之前训练好的模型出发,对模型进行优化,即模型首先从输入X(k)生成输出X(k+1),然后模型使用输出X(k+1)作为输入,并生成输出X(k+2)。通过将X(k+1)和X(k+2)与它们在训练数据中各自的真实值进行比较来计算训练损失,并使用这两个训练损失的总和来优化模型。整个训练过程是在一个由64个NvidiaA100GPU组成的集群上完成的,端到端的训练大约需要16个小时。降水模型ERA5再分析数据集中的总降水量(TP)是一个变量,表示通过降雨和降雪落到地球表面的累积液体和冻结水。该研究使用总降水量作为诊断变量并表示为p(kΔt)。用于训练主干模型的数据集(20个变量)不包括变量TotalPrecipitation。相反,我们训练一个单独的AFNO模型并使用骨干模型输出来诊断TP,如图2(c)所示。这种方法将降水建模的困难与预测大气状态的一般任务分离开来。此外,一旦经过训练,诊断TP模型可以与其他预测模型(传统的NWP或数据驱动的预测)相结合。用于从主干输出诊断降水的模型具有相同的基本AFNO架构,增加了额外的2D卷积层和ReLU激活以强制执行非负降水输出。由于主干模型以6小时为增量进行预测,因此本研究训练了诊断降水模型来预测6小时内累积的总降水量。实验结果下面的图1定性地展示了FourCastNet模型以0.25°-lat-long分辨率预测全球表面风速的能力。其中表面风速的大小是使用风速的纬度和经向分量计算的,即。下面的图1显示了使用FourCastNet提前96小时生成的说明性全球近地表风速预报。研究人员重点介绍了已解决和准确跟踪的预测的关键高分辨率细节,包括超级台风山竹(Mangkhut)(2018)和三个命名的气旋(佛罗伦萨、伊萨克和海伦)向美国东海岸移动。其中,图1(a)为模型初始化时的风速。图1(b)显示了模型提前96小时预报(上)和当时对应的真风速(下)。可以看出,FourCastNet模型能够提前96小时预测风速,具有出色的保真度和正确的精细尺度特征。值得注意的是,图1说明了名为“山竹”的超级台风的预测形成和轨迹,该台风在大约北纬10°、西经210°的初始位置开始形成。该模型定性地预测了台风的增强及其在4天内的轨迹以惊人的保真度。这表明FourCastNet模型擅长预测台风等天气现象的形成和轨迹。该研究使用样本外测试数据集的初始条件初始化FourCastNet模型。基于这个初始条件,模型在推理模式下可以自由运行16个时间步长,如下图2(d)所示。与台风相比,降水具有间歇性和随机性,因此预测降水量被认为是一项极其困难的工作。下面的图3展示了FourCastNet模型预测全球总降水量的能力。图3(a)为初始降水情况,图3(b)为模型提前36小时预测与实际情况对比。该图以美国西海岸局部地区的降水为例,突出显示了FourCastNet模型以非常高的精度解析和预测局部降水的能力。研究人员观察到,FourCastNet在捕获短时高分辨率降水特征方面表现良好,这对于预测极端天气事件具有重要意义。值得注意的是,这是深度学习模型首次成功用于大规模降水预测。接下来,研究人员进一步探索了开发深度学习模型来预测飓风的潜在效用。飓风是一种具有巨大破坏潜力的极端天气事件。建立快速可用且计算成本低廉的大气预警模型对于减轻生命和财产损失非常重要,并且此类模型还需要更严格的要求以避免错误预测。下面的图4展示了FourCastNet在预测快速变化的飓风方面的能力。除了台风和降水,FourCastNet模型还能够预测大气河流的形成和演化。下图5显示了2018年4月最终登陆北加州海岸线的一条大气河流的预报。此外,FourCastNet在量化技术上也有很大优势,优于IFS。对于ACC和RMSE指标,FourCastNet在更短的时间内(大约48小时或更短)实现了比IFS更好的性能。而FourCastNet只使用了IFS模型全变量集的一部分,计算成本也低很多。集合是数值天气预报的重要组成部分,FourCastNet使用大集合的性能提升如下图所示。感兴趣的读者可以阅读论文原文了解更多研究细节。