ECMWF的AI天气预报系统AIFS开放模型权重参数

文摘   2024-12-16 16:20   北京  

准确的天气预报对社会具有重要意义,有助于挽救生命和保护财产。近年来,数据驱动的天气预报领域取得了实质性进展(本号的另一篇推文AI天气模型的现状与未来)。继早期的探索性工作之后,谷歌、华为和英伟达等科技公司开发了数据驱动的天气预报模型,在多项标准预报评分(如500百帕位势高度均方根误差和异常相关系数)方面超越了领先的基于物理的全球数值天气预报模型。同时,这些预报可以在单个GPU(图形处理单元)上在几分钟内完成,相比之下,基于物理的数值天气预报需要数量级为1000的CPU。

这些数据驱动模型通常使用历史天气数据进行训练,主要使用欧洲中期天气预报中心(ECMWF)在欧盟哥白尼计划框架下生产的ERA5再分析数据集的子集。在进行实时预报时,这些模型目前依赖于业务数值天气预报的初始条件,如ECMWF的四维变分同化分析系统。

AIFS介绍

欧洲中期天气预报中心(ECMWF)开发了一个名为AIFS(人工智能预报系统)的数据驱动天气预报模型。这是一个可用于业务运行的实验性数据驱动预报系统,包括实现数据驱动预报模型,以及涵盖数据集生成、可重复训练、业务推理、业务检验工具使用、产品生成和预报数据分发的端到端流程。

AIFS这个名字是向ECMWF的IFS(集成预报系统)致敬。AIFS使用Python实现,利用PyTorch和PyTorch Geometric机器学习框架。与AIFS一起,研究团队还开发了Anemoi(希腊语:"风")框架,这是一个旨在为数据驱动天气预报模型提供构建模块的工具箱。

AIFS基于图神经网络(GNN)编码器和解码器,以及滑动窗口transformer处理器,通过ECMWF的ERA5再分析数据和业务数值天气预报(NWP)分析数据进行训练。AIFS具有灵活和模块化的设计,支持多层次并行计算,能够处理高分辨率输入数据。通过与NWP分析和直接观测数据比较,AIFS展现出对高空变量、地面天气参数和热带气旋路径的高度预报技能。目前AIFS每天运行四次,与ECMWF的物理模型并行运行,预报结果根据ECMWF的开放数据政策向公众开放。

AIFS系统架构

AIFS的第一个版本于2023年10月投入试验业务运行。其架构采用编码器-处理器-解码器设计。编码器将输入数据网格的信息聚合到低分辨率处理器网格上,解码器则将潜在状态映射回输入数据网格。在2024年2月,AIFS进行了升级,空间分辨率提升至约0.25度。编码器和解码器中的GNN改用了基于注意力/transformer的变体,处理器基于具有shifted window attention的pre-norm transformer。处理器网格基于O96八面体简化高斯网格,共有40320个网格点。AIFS总共有16层处理器层。

AIFS编码器/解码器和处理器框图示意图:GNN模块(左)、处理器模块(右)。GNN模块使用多头图转换器卷积操作来更新处理器的节点和边,而预归一化转换器模块则依赖于多头自注意力机制。

AIFS编码器(左)和解码器(右)图示例。ERA5网格点显示为小黑点,处理器网格点显示为较大的灰色圆圈。边界用黑线表示。为了清晰可见,仅展示了地球表面的一小部分区域。

不同格点(红色)的注意力窗口(蓝色高亮的格点)示意图。灰色格点展示了信息在6层处理器中能传递的最远距离(AIFS总共有16层处理器)。为了便于说明,图中显示的是比AIFS实际使用的分辨率更低的网格。

训练方法

AIFS被训练用于生成6小时预报。它接收t-6h和t0时刻的大气状态表示(来自ERA5或ECMWF业务分析)作为输入,然后预测t+6h时刻的状态。更长期预报通过自回归方式计算。训练分为两个阶段:预训练阶段针对6小时预报,随后是可以预报长达72小时的第二阶段。训练过程中使用了区域加权均方误差作为损失函数,并对不同输出变量应用了损失缩放。

预报效果

AIFS在北半球异常相关系数(ACC)等指标上表现出色,相比IFS在更长预报时效上具有超过12小时的预报优势。对流层顶以下的预报得分普遍优于IFS,改进幅度达到10%。地面预报验证结果也显示AIFS普遍优于IFS。在热带气旋路径预报方面,AIFS的位置误差明显低于IFS,这主要得益于减少了移动速度的慢偏差。

北半球500百帕位势高度距平相关系数(ACC)对比图,展示了AIFS(蓝色)、IFS(红色)和ERA5(黑色)在不同年份的预报效果:2天预报(左上)、6天预报(右上)和10天预报(底部)。

系统特点与局限

欧洲中期天气预报中心(ECMWF)开发了一个名为AIFS(人工智能预报系统)

AIFS具有高度灵活、可扩展的模块化架构。当前版本存在一些局限,如预报场在长期预报中会出现模糊化现象,这与均方误差训练目标有关。此外,在平流层预报技巧相对较低,热带气旋强度预报存在偏差。未来的改进方向包括发展概率集合预报、改进降水场表现、发展区域模式等。ECMWF计划在今年晚些时候以开源许可发布AIFS源代码和权重。

总的来说,AIFS代表了数据驱动天气预报的重要进展,展现出与传统物理模式相当甚至更优的预报能力,同时具有计算效率高、可扩展性强等优势。这个系统的发展对推动数值天气预报的创新具有重要意义。

这是AIFS与IFS(2022)预报评分的对比记分卡。预报初始场为00时和12时UTC。图中展示了不同预报时效(1至10天)在北半球外tropics区域(n.hem)、南半球外tropics区域(s.hem)、tropics区域和欧洲地区的相对评分变化。蓝色表示评分改进,红色表示评分降低。紫色表示预报距平标准差增加,绿色表示减少。带框的矩形表示达到95%显著性水平。


AIFS开放模型权重参数

之前,AIFS已经开源了其用于开发AI天气预报模型的框架Anemoi的代码。近期,最新的确定性版本AIFS v0.2.1(AIFS-single)的权重参数已在Hugging Face平台上发布。这使得开发者、科学家和天气爱好者都能在自己的设备上运行AIFS模型。该模型采用知识共享署名4.0国际许可协议(CC BY 4.0)发布,这意味着研究人员和实践者可以以任何方式使用和试验该模型。

在Hugging Face平台上,AIFS模型权重附带详细的"模型卡",其中包含模型实现和评估的具体细节。为了演示如何生成预报,模型卡中提供了可在Google Colab等系统上运行的notebook链接,其中包含详细的设置步骤指南和代码片段。

使用AIFS v0.2.1的主要步骤包括:

  1. 安装所需软件包

  2. 从ECMWF开放数据获取初始条件

  3. 使用anemoi-inference生成预报

  4. 检查生成的预报结果


该系统提供了多种应用程序接口(API),包括用于生成预报的NumPy-to-NumPy API,以及用于自定义工作流的命令行接口,为不同需求的用户提供了灵活的使用方式。

这一进展标志着AIFS在推动开放科学和促进气象预报领域协作方面迈出了重要一步,为未来更广泛的应用和创新奠定了基础。

References:

  • 开放数据集:

    https://www.ecmwf.int/en/forecasts/datasets/open-data

  • 模型权重下载:

    https://huggingface.co/ecmwf/aifs-single

  • Anemoi框架开放代码:

    https://github.com/ecmwf/anemoi-inference

  • 论文:https://arxiv.org/abs/2406.01465


本公众号相关内容推荐


气python风雨
主要发一些涉及大气科学的Python文章与个人学习备忘录
 最新文章