实验｜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更多干货奉上

幻方 | 技术博客

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

探索AI实践最优解，AISummit全球人工智能技术大会完美落幕
探索AI实践最优解,AISummit全球人工智能技术大会完美落幕北京时间2022年8月7日下午17:30,由51CTO精心策划以"驱动•创新•数智"为主题的AISummit全球人 ...
AI驱动的超分辨技术落地实践
近年来,随着深度学习技术的快速发展,基于AI的超分辨技术在图像恢复和图像增强领域呈现出广阔的应用前景,受到了学术界和工业界的关注和重视.但是,在RTC视频领域中,很多AI算法并不能满足实际场景下的应用 ...
Kyndryl从IBM完全剥离在纽交所独立上市；新思科技收购AI驱动性能优化软件企业Concertio | 全球TMT...
海外市场 Kyndryl (勤达睿)宣布已经完成从 IBM的剥离,使用股票代码"KD"在纽约证券交易所开始交易,Kyndryl 在纽约证券交易所敲响了开市钟.Kyndryl 于 1 ...
模型实践 | 高精地图构建模型HDMapNet助力更精准的自动驾驶
实验 | Freja 算力支持 | 幻方AIHPC 高精地图是自动驾驶系统的关键模块,可以有效提升自动驾驶汽车的行驶安全度,强化自动驾驶系统的整体感知能力和决策能力.然而传统的高精地图构建流程复杂 ...
AI驱动洞察智能数据营销时代迎来拐点
「在美国,科技业的一半是to C的公司,他们占了一半的市值,to B的这些公司占了另外一半市值,比如说Oracle.所以其实to B还有很多很牛逼的公司,但是我们来看中国的话发现中国to C的公司都很 ...
聚焦 AI + 大数据全球视野引领行业创新升级
10 月 17 日,神策 2017 数据驱动大会在北京丽亭华苑酒店如期举行.大会为期两天,来自全球大数据和人工智能领域的领袖人物聚首北京,围绕"智能·追本数源"主题,从全球视野与中 ...
通过客户流失预测案例感悟数据分析设计方法思考——数据驱动、AI驱动
国际著名的咨询公司Gartner在2013年总结出了一套数据分析的框架,数据分析的四个层次:描述性分析.诊断性分析.预测性分析.处方性分析. Gartner于2020年中给出预测,到2024年底,75 ...
AI驱动的智能图形应用
高速5G和性能更强大的终端为社交.游戏.教育等场景带来的更丰富的互动体验,本文内容来自相芯科技的蔡锐涛在LiveVideoStackCon2019深圳站上的精彩分享,他将介绍AI驱动的智能图形应用,以 ...
【金猿产品展】Stratifyd——敏捷AI驱动的客户体验数据分析平台
Stratifyd产品本项目由Stratifyd投递并参与"数据猿年度金猿策划活动--2021大数据产业创新服务产品榜单及奖项"评选. 数据智能产业创新服务媒体 --聚焦数智 · ...

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

数据集