模型实践 | AI驱动的全球精准气象预报

实验｜Vachel       算力支持｜幻方AIHPC 

现代数值天气预报（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

---

数据集

欧洲中期天气预报中心（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 等功能。

1►训练样本集

如前面所述，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 文档：https://doc.hfai.high-flyer.cn/index.html

2►模型训练

训练分为 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 的复现与优化，使得该项研究成果切实落地，为环境治理、灾害预警贡献一份力量。

综合体验打分如下：

01：研究新颖度  ★★★★★

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

02：开源指数  ★

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

03：算力门槛  ★★★★

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

04：通用指数  ★★★

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

05：模型适配度  ★★★★★

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

 

幻方AI BLOG更 多 干 货 奉 上

幻方 | 技术博客