FourCastNet3 (FCN3) 是 NVIDIA Earth-2 的最新 AI 全球天气预报系统。FCN3 在亚季节时间尺度上提供了概率技能、计算效率、频谱保真度、集成校准和稳定性的前所未有的组合。其中期预报准确度与 GenCast 等领先的机器学习模型相媲美，并超过 IFS-ENS 等传统数值天气预报系统。
在单个 NVIDIA H100 Tensor Core GPU 上，一个 60 天的 FCN3 0.25°、6 小时分辨率的部署在不到 4 分钟内完成，速度比 GenCast 快 8 倍，比 IFS-ENS 快 60 倍。
它还具有出色的校准和光谱保真度，即使在 60 天的延长预报时间内，集合成员也能保持逼真的光谱特性。FCN3 在从中期到亚季节时间尺度的较大集合数据驱动天气预报方面实现了重大飞跃。
FCN3 架构
FourCastNet3 采用基于球形信号处理原语的全卷积球形神经运算符架构（见图 2）。与基于球面傅立叶神经运算符的 FourCastNet2 不同，FCN3 在光谱卷积的同时还使用了局部球面卷积。
这些卷积使用 Morlet 小波进行参数化，并在离散连续群卷积框架中进行制定。这种方法能够实现各向异性局部滤波器，非常适合局部大气现象，同时通过 NVIDIA CUDA 中的自定义实现保证计算效率。
FCN3 通过一个潜在噪声变量在每个预测步骤中引入随机性，该变量的演化受球面扩散过程的控制。这种隐马尔可夫公式能够高效地一步生成集成成员，这是基于扩散模型的方法的一大优势。FCN3 作为集成模型进行联合训练，最大限度地减少了空间和频谱域中连续排名概率分数 (CRPS) 的复合损失函数。这种方法可确保 FCN3 能够学习到潜在随机大气过程中的正确空间相关性。
扩展机器学习模型通常是获得竞争优势的关键，但数据驱动型天气模型的扩展效果尚未得到研究。FCN3 的计算目标不同寻常。为了扩展它，我们引入了一种新的模型并行范例，该范例的灵感来自传统数值天气建模中的域分解。
这种方法使我们能够通过在多个设备上拆分模型，在训练期间将更大的模型放入 VRAM，同时降低每个设备的磁盘 I/O。为此，使用 NVIDIA 集群通信库 (NCCL) 以分布式方式实现卷积等空间操作。使用这项技术，FCN3 在多达 1024 个 GPU 上进行训练，同时使用域、批次和集成并行。查看我们的训练代码。
FourCastNet3 的表现优于基于物理的最佳集成模型 IFS-ENS，在预测能力方面与 Gencast 相当（见图 3）。在单个 NVIDIA H100 上，FCN3 能够在一分钟内以 6 小时时间分辨率和 0.25° 空间分辨率生成 15 天的单一预测，速度比 Gencast 快 8 倍，比 IFS-ENS 快 60 倍。
其概率集合的传播技能比始终接近 1，表明预测的不确定性与观测到的大气变化密切相关，从而实现了校准良好的预测。排名直方图和额外的诊断结果证实，集成成员与现实世界观察结果仍然可以互换，从而确认了 FCN3 预测的可靠性和可信度。
至关重要的是，FCN3 保留了所有尺度的大气光谱特征，即使在长达 60 天的预报期内，也能忠实地再现现实世界天气模式的能量级联和清晰度。与许多 ML 模型不同，FCN3 不会模糊高频特征或随着时间的流逝而退化成噪声伪影，而是保持稳定、物理上逼真的光谱，从而在亚季节范围内提供准确、清晰且物理上一致的预测。
图 3 显示了 FCN3 在 2020 年 2 月 11 日（风暴丹尼斯登陆欧洲前不久）对 500 毫巴风速的预测。FCN3 可准确捕捉不同长度尺度上风强度的幅度及其变化，这通过各自预测的忠实角度功率谱密度得以说明。即使在 30 天（720 小时）或更长时间的扩展部署中，情况也是如此。
开始使用 FourCastNet3
经过全面训练的 FourCastNet3 检查点可在 NVIDIA NGC 上获得。
运行 FCN3 推理的一种简单方法是使用 Earth2Studio。要运行单个 4 个成员的集成推理，您可以执行以下代码：
from earth2studio.models.px import FCN3
from earth2studio.data import NCAR_ERA5
from earth2studio.io import NetCDF4Backend
from earth2studio.perturbation import Zero
from earth2studio.run import ensemble as run
import numpy as np
# load default package
model = FCN3.load_model(FCN3.load_default_package())
# determine output variables
out_vars = ["u10m", "v10m", "t2m", "msl", "tcwv"]
# data source initial condition
ds = NCAR_ERA5()
io = NetCDF4Backend("fcn3_ensemble.nc", backend_kwargs={"mode": "w"})
# no perturbation required due to hidden Markov formulation of FCN3
perturbation = Zero()
# invoke inference with 4 ensemble members
run(time=["2024-09-24"],
nsteps=16,
nensemble=4,
prognostic=model,
data=ds,
io=io,
perturbation=perturbation,
batch_size=1,
output_coords={"variable": np.array(out_vars)},
)
图 4 显示了此推理的结果。为了获得最佳 FCN3 性能，我们建议安装 torch-harmonics，同时启用自定义 CUDA 扩展，并在推理过程中使用 bf16 格式的自动混合精度（这是 Earth2Studio 的默认设置）。如果您想运行自定义 FCN3 推理或自己训练 FCN3，可以在 makani 中找到代码。
详细了解 FCN3
通过以下资源详细了解 FourCastNet3：
- FourCastNet3 论文arXiv。
- FourCastNet3 检查点
- GitHub 上的 FCN3 分布式训练/推理代码。
- 使用torch-harmonics进行球形信号处理。
- Earth2Studio用于推理。
- Physicsnemo：用于物理模拟的深度学习框架。
- arXiv 上的球面傅立叶神经运算符论文
- 使用球面傅立叶神经运算符对地球大气进行建模
- FourCastNet2 (SFNO) 检查点
完整作者列表
Boris Bonev (NVIDIA) 、Thorsten Kurth (NVIDIA) 、Ankur Mahesh (LBNL) 、Mauro Bisson (NVIDIA) 、Karthik Kashinath (NVIDIA) 、Anima Anandkumar (加州理工学院) 、William D. Collins (LBNL) 、Mike Pritchard (NVIDIA) 、Alex Keller (NVIDIA)
