超算跑模型 | AI驱动的全球精准气象预报

High-Flyer    June 10, 2022

现代数值天气预报(numerical weather prediction, NWP)可以追溯到1920年,其基于物理原理,整合了几代气象学者的成果经验,是各国气象部门所采用的主流的天气预报方法。这其中,来自欧洲中期天气预报中心(ECMWF)的高分辨率综合预测系统(IFS)模型,在历年的全球天气预报中表现最佳。

最近,在刚刚闭幕不久的Nvdia GTC技术年会上,来自英伟达、劳伦斯伯克利国家实验室、密歇根大学安娜堡分校、莱斯大学等机构的研究者研发了一种基于傅里叶神经网络的预测模型 FourCastNet,它能以 0.25° 的分辨率生成全球关键天气指标的预测,这相当于赤道附近大约 30×30 km 的空间分辨率和 720×1440 像素的全球网格大小,与IFS系统一致。这项成果使得AI气象模型首次能够与传统物理模型IFS进行直接比较

FourCastNet 采用GPU训练计算,比传统的 NWP 模型快约 45000 倍,能量节约12000倍,使得它能够以很低的成本生成大量的子模型,进行集合预测。实验表明 FourCastNet 极大地改善了概率天气预报的效果,可以在几秒钟内生成对飓风、大气层河流和极端降水等事件的大规模集合预报。

幻方AI最近复现了该项工作,并通过幻方自研的 3FShfreduce算子,对模型训练和推导进行优化。我们在hfai数据仓库中开源了训练数据,模型代码,旨在帮助研究者和开发者们降低研究门槛。您可以在幻方萤火超算集群中快速体验,申请试用

论文标题:FourCastNet: A Global Data-driven High-resolution Weather Model using Adaptive Fourier Neural Operators

论文地址https://arxiv.org/abs/2202.11214

模型仓库https://github.com/HFAiLab/FourCastNet

数据集

欧洲中期天气预报中心(ECMWF)提供了一个公开可用的综合数据集 ERA5。作为第五代全球气候再分析结果,ERA5将其物理模型数据与来自世界各地的观测数据结合起来,形成一个全球完整的、一致的数据集,以小时级到天级不等,提供包括温度、风量、降水、水文、气压等多项全球气象指标数据,供各种气象预报模型学习。官方地址如下:https://www.ecmwf.int/en/forecasts/datasets/reanalysis-datasets/era5

FourCastNet使用 ERA5 来训练,专注于预报两大气象变量:距离地球表面10m处的风速6小时总降水量,以对极端天气、自然灾害的预警。该两项指标的预测也被认为是气象领域比较难准确预测的问题之一。

为此,FourCastNet选择了20个相关气象指标,包括几个不同垂直高度的位势高度、温度、风速和相对湿度,一些近地表变量,如地面气压和平均海平面气压以等,作为模型的输入,以挖掘出变量间彼此的关系对风速与降水的影响。具体变量如下:

0.png

模型介绍

为了进行0.25° 分辨率下的全球气象预测,FourCastNet采用自适应傅里叶神经算子 AFNO,这种神经网络架构是对Vision Transformer(ViT)模型地改进,它将混合操作步骤构建成连续的全局卷积,在傅里叶域中通过 FFT 有效实现。通过这样的设计,将空间混合复杂度降低到 O(N log N),这允许灵活且可扩展地建模跨空间和通道维度的依赖关系。AFNO的结构如下图所示:

1.png

在空间混合步骤中,GFNet使用跨通道的逐元素矩阵乘法,FNO采用混合所有通道的全矩阵乘法,而AFNO结合MLP和软阈值进行逐块矩阵乘法。

对于气象数据,该研究将多种气象变脸按时间整合成 [time, variable, latitude, longitide]格式,进行AFNO。具体的:FourCastNet 首先将 720 × 1440 lat-lon 网格上的输入变量投影到 2D patch 网格 (h × w)(patch大小为 p × p,例如 p = 8),每个 patch 表示为一个 d 维 token。然后,将 patch 序列连同位置编码一起馈送到一系列 AFNO 层。如下所示:

2.png

FourCastNet的训练采用递进式,即以 X(t) 作为输入,预测下一步X(t+1)。一次训练输出多步,与真值对比计算loss。这里论文里采用RMSE作为loss的评价指标。

对于降水模型,如上图(c)所示,每一步降水的预测实在主干AFNO输出的基础上套接一个降水专有AFNO进行综合判断,其基于主干网络的预训练结果。这是因为降水这一指标与其他气象指标的概率分布不同,呈现比较鲜明的长尾特征,使得特征空间非常稀疏。对于降水专有的AFNO,其与主干AFNO基本一致,不同在于输出侧,使用Relu进行激活,以防止非负输出。

模型复现

幻方 AI 基于AFNO和FourcastNet两篇论文的描述,复现了 FourCastNet,并采用幻方一系列优化工具进行提速升级,包括ffrecordhfreducehfai.datasetshfai.checkpoints等功能。

训练样本集

如前面所述,FourCastNet模型的训练包括pretrianfinetuneprecipitation 三个部分,对应的数据输入各有不同。因此,我们将原始的 ERA5 数据进行清洗,构建如下样本数据集:

for xt, xt1, xt2, pt1 in hfai.datasets.ERA5.loader():
    out1 = backbone_model(xt)
    out2 = backbone_model(out1)
    out3 = precip_model(out1)
    
    loss_pretrain = criterion(xt1, out1)
    loss_finetune = criterion(xt2, out2) + loss_pretrain
    loss_precip = criterion(pt1, out3)
    
    ...

其中xt代表当前t时刻的全球20各气象指标数据,xt1代表t+1时刻,xt2代表t+2时刻,pt1代表t+1时刻的降水指标数据。依次进行pretrian、finetune 和 precipitation的训练。

目前该样本数据集已整合进幻方数据集仓库中,可以通过hfai.datasets.ERA5调取使用。更多参考hfai文档

模型训练

训练分为pretrian、finetune 和 precipitation 三个部分,每个部分的输入输出各有不同。

pretrain部分,以 X(t) 作为输入,预测下一步X(t+1),旨在缩小预测的误差。训练80个轮次。

for xt, xt1, _, _ in hfai.datasets.ERA5.loader():
    out = backbone_model(xt)
    loss_pretrain = criterion(xt1, out)
    
    ...

finetune部分,在原有主干模型(backbone)的基础上,多预测一步X(t+2),进一步提升模型预报的精度。训练50个轮次。

for xt, xt1, xt2, _ in hfai.datasets.ERA5.loader():
    out1 = backbone_model(xt)
    out2 = backbone_model(out1)
    loss_finetune = criterion(xt1, out1) + criterion(xt2, out2)
    
    ...

preciptation部分,原有的主干模型训练完毕,以X(t+1)的信息,预测同时刻的降水P(t+1)。训练25个轮次。

for xt, _, _, pt1 in hfai.datasets.ERA5.loader():
    out = backbone_model(xt)
    out_p = precip_model(out)
    loss_precip = criterion(pt1, out_p)
    
    ...

我们采用hfai.ddp调用hfreduce加速训练,使用64张A100进行数据并行加速,数据和模型采用半精度,最终整体训练耗时在13~14小时左右。

更多内容可以阅读我们模型仓库中的开源:https://github.com/HFAiLab/FourCastNet

训练结果

我们以2018年9月的特大型台风“山竹”为例,展示我们复现的效果:

台风路径

3.gif

总水汽浓度(Total column water vapour)

4.gif

可以看到,模型准确预测了未来3天的台风行踪路径,风力和降水范围,强度。

体验总结

FourCastNet 将AFNO应用于气象领域,首次将AI气象的分辨率提升至与物理模型同等水平,使得AI气象模型应用于真实场景中可能性大幅提高。FourCastNet在大风、极端暴雨的预测表现,为进一步提升天气预报的准确度,预防极端自然灾害提供了新的方案。幻方AI对FourCastNet的复现与优化,使得该项研究成果切实落地,为环境治理、灾害预警贡献一份力量。

综合体验打分如下:

  1. 研究指数:★★★★★

    该模型首次将AI气象分辨率和准确度提升至物理模型水平,促进了AI气象模型的落地。同时,与气象领域知识的融合让该方向有更多研究空间。

  2. 开源指数:★

    代码没有开源,ERA5数据下载比较繁琐。

  3. 门槛指数:★★★★

    数据规模大,模型适中,适合多级多卡数据并行训练。一般单卡训练难度比较大。

  4. 通用指数:★★★

    该方法是基于AFNO的应用,在高分辨率数据场景中皆可以使用。

  5. 适配指数:★★★★★

    依赖简单,很容易与幻方AI的训练优化工具结合,提效明显。

幻方 AI Lab 紧跟 AI 研究的前沿浪潮,致力于用领先算力助力AI落地与价值创造,降低AI的使用门槛,欢迎各方数据科学家与开发者们一同共建。


本文作者: High-Flyer


您可以转载、不违背作品原意地摘录及引用本技术博客的内容,但必须遵守以下条款: 署名 — 您应当署名原作者,但不得以任何方式暗示幻方为您背书,亦不会对幻方的权利造成任何负面影响。 非商业性使用 — 您不得将本技术博客内容用于商业目的。 禁止演绎 — 如果基于该内容改编、转换、或者再创作,您不得公开或分发被修改内容,该内容仅可供个人使用。