DeepMind 如何利用 AI 和 JAX 预测气旋

传统上，提前一周预测天气需要依赖强大的超级计算机。这些计算机需要花费数小时在庞大的网格上求解流体力学和热力学方程。Google DeepMind 采用了一种不同的方法：工程师们训练神经网络在几秒钟内预测大气过程。该仓库 google-deepmind/weathernext 包含 WeatherNext 2、GraphCast 和 GenCast 架构的源代码和权重，其中包括专门用于追踪热带气旋的模型。

让我们来探索这个项目的结构、运行它所需的硬件资源，以及开发者或研究人员可以从中获取什么。

仓库里有什么

该项目整合了 DeepMind 几代气象神经网络的成果。重点是 WeatherNext 2（WN2）——一个全球中期天气预报模型。它的空间分辨率为 0.25 度（网格步长约为地球表面 30 km）。

该模型接收大气的初始状态并进行自回归运行，逐步生成温度、不同高度的 风速、位势高度以及气旋移动轨迹的预报。

仓库中包含三类模型：

• WeatherNext 2（<2025）：主力模型。它在业务化的 ECMWF HRES（高分辨率预报）数据上进行了微调，因此可以直接从真实天气分析中初始化。除了标准指标外，它还能预报 100 米高度的风速，这对风力发电至关重要。
• WeatherNext Cyclones：一系列针对飓风和热带风暴追踪优化的模型（2023、2024 和 2025 版本）。2025 版本称为 FNV3，已在真实的大西洋飓风季中经过了实战检验。
• WeatherNext Cyclones Mini：分辨率为 1 度的轻量级模型。它们小得多，适合本地测试或在单块 GPU 上运行。

除了 WN2，文件夹 docs/ 还包含早期模型的链接和资料：基于图网络的 GraphCast 和基于扩散的 GenCast。

硬件和系统要求

以 0.25° 全分辨率运行完整模型需要大量资源。重型检查点的推理针对 TPU（特别是 TPU v5p）进行了优化。如果你想在 GPU 上运行完整版本，由于高分辨率张量运行所需的内存量，你需要一块 NVIDIA H100 级别的显卡。

为了便于探索和实验，创作者发布了 Mini 版本。这个版本对开发者很友好：

• 可以在免费的 Google Colab 上流畅运行，使用 TPU v5e-1。
• 即使在较老的加速器（如 NVIDIA P100 或 T4）上也能运行。

代码结构及入门指南

代码使用 Python 编写，基于 JAX 和库 xarray 处理多维气象数据数组。

安装该库只需运行：

pip install git+https://github.com/google-deepmind/[email protected]

作者强烈建议锁定特定版本，因为这是研究代码，API 稳定性无法保证。

通过 Colab 快速入门

文件夹 docs/weathernext2/wn2_demo.ipynb 包含一个即用型笔记本。它演示了使用该模型的完整工作流程：

1. 从公共 Google Cloud 存储桶加载权重（dm_graphcast）。
2. 准备初始大气条件（通过 Zarr 格式和 WeatherBench2 平台使用 ERA5 或 HRES 数据）。
3. 在 JAX 中初始化 FGN/WN2 架构。
4. 运行自回归展开以生成未来数天的预报。
5. 使用内置追踪器计算气旋轨迹。
6. 计算损失并执行一次梯度步（是的，你可以本地微调模型）。

以下代码中概念性地展示了预报循环的样子：

# Загрузка весов и конфигурации
model_weights = load_weights("WeatherNextCyclones_Mini_<2024.npz")
initial_state = load_hres_analysis(timestamp="2025-01-01T00:00")

# Авторегрессионный прогон на несколько шагов вперед
forecast_steps = 14  # например, прогноз на 3.5 дня с шагом 6 часов
predictions = run_autoregressive_rollout(
    model=wn2_architecture,
    params=model_weights,
    inputs=initial_state,
    steps=forecast_steps
)

实际应用

如果你不打算在自己的天气数据上微调模型或构建科学流程，那么设置 JAX 基础设施并非必要。Google 向开放源发布了现成的 WN2 结果：

• OpenMeteo API：你可以通过常规 REST 请求获取 WN2 数据，无需连接 GPU。
• Google Cloud / BigQuery / Vertex AI：用于分析和大规模查询的现成数据集。
• WeatherLab：用于可视化飓风轨迹的交互式服务。

对于从事可再生能源研究、气候风险分析或构建自己的 ML 服务的人来说，仓库中的代码是一个现成的基础。第 utils/ 节包含便捷的数据归一化工具、图块处理和特定气象损失计算功能。

总结

WeatherNext 是深度学习如何逐步取代或补充传统数值天气预报（NWP）的一个出色示范。

优点：

• 开放的权重和可在 Colab 中运行的演示笔记本。
• 轻量级的 Mini 模型，可在常规硬件上进行实验。
• 与真实气象数据（WeatherBench2、ERA5、HRES）的现成集成。

缺点：

• 全分辨率推理需要顶级的 H100 或 TPU 加速器。
• API 不保证版本间的向后兼容性。

如果你使用 JAX 或想了解如何在大型张量网格上运行自回归传递，一定要试试 WN2 演示笔记本。