【时空序列】AI驱动的全球精准气象预报

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

AI气象预测

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

推荐阅读：
 
我的2022届互联网校招分享
 
我的2021总结
 
浅谈算法岗和开发岗的区别
 
互联网校招研发薪资汇总
2022届互联网求职现状，金9银10快变成铜9铁10！！
 
公众号：AI蜗牛车
 
保持谦逊、保持自律、保持进步
 
发送【蜗牛】获取一份《手把手AI项目》（AI蜗牛车著）
发送【1222】获取一份不错的leetcode刷题笔记
 
发送【AI四大名著】获取四本经典AI电子书