GraphCast 如何用图神经网络把全球中程天气预报做到 0.25° 精度
【免费下载链接】weathernext项目地址: https://gitcode.com/GitHub_Trending/gr/weathernext
GraphCast 是 Google DeepMind 开源的全球中程天气预测模型,可在 0.25° 分辨率上逐步推演温度、风场、位势高度等变量,单次前向 12 小时一步,滚动外推约 10 天;同库还提供扩散模型版 GenCast。
GraphCast 模型规格与预测能力边界
该项目提供两套模型。GraphCast 为确定性点预测,共三个规格:
- GraphCast:0.25° 分辨率、37 个气压层,用 ERA5 数据(1979–2017 年)训练,对应论文主模型;
- GraphCast_small:1° 分辨率、13 个气压层,显存与算力需求更低,便于本地验证;
- GraphCast_operational:0.25° 分辨率、13 个气压层,ERA5 预训练后在 HRES 数据(2016–2021 年)上微调,可由 HRES 场初始化,且不依赖降水输入,面向业务化场景。
GenCast 为概率性集合预报,提供四个规格:两个 0.25° 模型(其中 Operational 版经 HRES-fc0 微调)、一个 1° 模型,以及一个1° Mini模型——它显存占用最小,可在免费 Colab 中直接运行,但精度不代表其余三个模型的水平。
预测时长以中程为主,评估分数卡中预报时长(lead time)最多延伸到 15 天。其能力边界需明确:GraphCast 输出单一确定场,GenCast 输出概率集合;两者都不替代官方气象警报,且输入依赖 ERA5、HRES 等第三方再分析数据,受其各自授权条款约束。
上图为 GenCast 1° Mini 相对 ENS(集合预报系统)的 RMSE 与 CRPS 分数卡,纵轴为气压层、横轴为预报时长:蓝色代表优于参照系统,红色代表落后,可见低空变量随预报时长增长精度衰减更快。
三段式图神经网络架构与 BFloat16 精度设计
网格与球面三角网格之间的消息传递
数据以经纬度格点组织,而模型内部计算放在球面三角网格上,中间用三段图神经网络衔接:
- Grid2Mesh:以半径查询为格点与三角网格节点建边,单步消息传递把格点特征送入网格;
- Mesh:在合并的多层网格上执行多步消息传递,让不同尺度的空间信息相互传播;
- Mesh2Grid:把落在每个三角面内的网格点回插出最终预测。
这样设计的原因在于:直接对全球格点做全连接代价过高,而球面网格让信息只在邻近节点间传递,多网格结构则同时覆盖粗、细两种尺度。节点与边用TypedGraph按类型区分管理,经纬度、相对位置等空间量作为结构特征参与计算;损失按气压层加权并对纬度做加权,使高纬区域不被低纬的密集格点淹没。
BFloat16 训练与 GenCast 的扩散采样
GraphCast 通过 BFloat16Cast 包装器以 BFloat16 精度运行——这是一种 16 位浮点格式,参数仍以 float32 存储,仅激活与计算走 BFloat16,梯度继续以 float32 累积,用约一半显存换取几乎无损的精度,对资源受限环境尤为关键。训练阶段用可微的自回归模块(把上一步输出喂回下一步)拟合序列,推理阶段改用 Python 循环的 rollout 模块推更长轨迹。
GenCast 在去噪框架上工作:一个 denoiser 负责单步去噪,采样器独立可配,默认使用 DPM-Solver++ 2S 快速采样器;训练时按噪声水平对损失加权,从而学到天气场的概率分布而非单点值。
从 ERA5 数据到可运行推理的上手路径
模型权重、标准化统计量与示例输入都放在 Google Cloud Bucket(dm_graphcast)的graphcast/与gencast/子目录中,按模型名对应。完整训练需要 ERA5 数据(来自 ECMWF),建议用 Weatherbench2 提供的 Zarr 格式访问;业务化微调则需 HRES-fc0 数据。
三个关键入口:
- 最低门槛:打开 gencast_mini_demo.ipynb,可在免费 Colab 中加载 Mini 权重(或随机权重)、生成预测、计算损失与梯度,无需 GPU;
- GraphCast 同理,入口为
graphcast_demo.ipynb,单步架构见graphcast.py; - 高分辨率模型(0.25°)需要 TPU,docs/cloud_vm_setup.md 指导启动 Google Cloud TPU VM 并通过
gencast_demo_cloud_vm.ipynb运行,GPU 方案需替换注意力实现。
上图为 GenCast 0.25° 模型在 H100 与 TPUv4 上的精度对比,说明不同加速器对结果存在可测量的差异,选型时值得实测。
适用人群与典型场景
- 算法研究者:复现或改造 GNN 架构,研究消息传递步数、多网格层数对精度的影响;
- 资源受限的演示与教学:用 1° Mini 模型在免费 Colab 完成一次完整推理与梯度演示;
- 业务化实验:以 operational 模型从 HRES 场初始化,评估概率集合预报的 CRPS 表现。
许可证与学术引用
代码遵循 Apache License 2.0;模型权重遵循 CC BY-NC-SA 4.0,仅限非商业用途。使用时建议引用两篇论文:
@article{lam2023learning, title={Learning skillful medium-range global weather forecasting}, author={Lam, Remi and Sanchez-Gonzalez, Alvaro and Willson, Matthew and Wirnsberger, Peter and Fortunato, Meire and Alet, Ferran and Ravuri, Suman and Ewalds, Timo and Eaton-Rosen, Zach and Hu, Weihua and others}, journal={Science}, volume={382}, number={6677}, pages={1416--1421}, year={2023}, publisher={American Association for the Advancement of Science} }@article{price2023gencast, title={GenCast: Diffusion-based ensemble forecasting for medium-range weather}, author={Price, Ilan and Sanchez-Gonzalez, Alvaro and Alet, Ferran and Andersson, Tom R and El-Kadi, Andrew and Masters, Dominic and Ewalds, Timo and Stott, Jacklynn and Mohamed, Shakir and Battaglia, Peter and Lam, Remi and Willson, Matthew}, journal={arXiv preprint arXiv:2312.15796}, year={2023} }这套开源实现把 GNN 与扩散模型两条技术路线同时落到可复现的代码与权重上,为后续在网格化地球科学任务上的架构研究提供了一个可直接对照的基准。
【免费下载链接】weathernext项目地址: https://gitcode.com/GitHub_Trending/gr/weathernext
创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考