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

Vachel 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 最近复现了该项工作，并通过幻方自研的 3FS、hfreduce、算子，对模型训练和推导进行优化。我们在 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

演示Demo：https://www.high-flyer.cn/hf-earth/

数据集

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

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

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

模型介绍

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

在空间混合步骤中，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 层。如下所示：

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

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

模型复现

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

训练样本集

如前面所述，FourCastNet 模型的训练包括pretrian、finetune 和 precipitation 三个部分，对应的数据输入各有不同。因此，我们将原始的 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 月的特大型台风“山竹”为例，展示我们复现的效果：

台风路径

总水汽浓度（Total column water vapour)

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

体验总结

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

综合体验打分如下：

研究指数：★★★★★

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

代码没有开源，ERA5 数据下载比较繁琐。
门槛指数：★★★★

数据规模大，模型适中，适合多级多卡数据并行训练。一般单卡训练难度比较大。
通用指数：★★★

该方法是基于 AFNO 的应用，在高分辨率数据场景中皆可以使用。
适配指数：★★★★★

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

幻方 AI 紧跟 AI 研究的前沿浪潮，致力于用领先算力助力 AI 落地与价值创造，欢迎各方数据研究者与开发者们一同共建。