炸裂！英伟达AI天气预报提速60倍，传统预报要凉？

FourCastNet3 (FCN3)是英伟达（NVIDIA）Earth-2推出的最新一代全球人工智能天气预报系统。它在概率预测能力、计算效率、频谱保真度、集合校准和亚季节尺度稳定性等方面实现了前所未有的结合。其在中等时间范围内的预测精度与GenCast等领先的机器学习模型相当，甚至超过了传统的数值天气预报系统，如IFS-ENS。

在一台英伟达H100 Tensor Core GPU上，FCN3仅用不到四分钟即可完成一次60天的模拟，分辨率为0.25°，每6小时输出一次结果，速度比GenCast快8倍，比IFS-ENS快60倍。此外，它还具有出色的校准和频谱保真度，即使在长达60天的预测中，集合成员也能保持真实的频谱特性。FCN3的出现，标志着数据驱动的天气预测技术向中长期亚季节尺度的大集合预测方向迈出了重要一步。

图为地球的动态展示。

下图展示了15个FourCastNet3集合成员在两周内的模拟结果，显示了这段时间内的地表风速变化情况。

图1：15个FourCastNet3集合成员在两周内的模拟结果，显示了这段时间内的地表风速。

FCN3的架构

FCN3采用了一种完全卷积的球形神经算子架构，其基础是球形信号处理原语（如图2所示）。与基于球谐傅里叶神经算子的FourCastNet2不同，FCN3同时使用了局部球形卷积和频谱卷积。这些卷积使用莫莱小波进行参数化，并在离散-连续群卷积框架中进行公式化。这种方法能够实现各向异性的局部滤波器，非常适合局部的气候现象，并通过英伟达CUDA中的自定义实现来保证计算效率。
图片说明

图2：FCN3架构图。FCN3是一种用于球形信号的神经算子，可以将当前时间步长的大气和地表变量映射到下一个时间步长。随机性是通过隐藏的马尔可夫模型方法引入的，该方法将球形噪声变量作为条件输入。

FCN3通过潜在噪声变量在每个预测步骤中引入随机性，该变量的演变由球体上的扩散过程控制。这种隐藏马尔可夫公式能够有效地单步生成集合成员，这是优于基于扩散模型方法的一个关键优势。FCN3作为一个整体进行联合训练，最小化一个组合损失函数，该函数结合了空间和频谱域中的连续排序概率分数（CRPS）。这种方法确保了FCN3能够学习到潜在随机大气过程中正确的空间相关性。
图片说明

图3：FCN3在2020年样本外验证年中，以12小时为间隔计算的概率分数。从上到下：连续排序概率分数（CRPS）、集合平均均方根误差（RMSE）、离散技巧比（SSR）和等级直方图。

扩展机器学习模型通常对于获得竞争优势至关重要，但规模效应尚未在数据驱动的天气模型中得到充分研究。FCN3的不同寻常之处在于其计算目标。为了扩展它，我们引入了一种新颖的模型并行范例，其灵感来自传统数值天气模型中的领域分解。这种方法使我们能够通过在多个设备上拆分模型，从而将更大的模型放入训练期间的VRAM中，同时降低每个设备的磁盘I/O。为此，使用英伟达集体通信库（NCCL）以分布式方式实现卷积等空间操作。利用这项技术，FCN3可以在多达1,024个GPU上进行训练，同时使用领域、批次和集合并行性。

外媒报道称，FCN3在预测能力方面优于最好的基于物理的集合模型IFS-ENS，并且与Gencast相匹配（见图3）。在新媒网跨境获悉，在单个英伟达H100上，FCN3在一分钟内生成单个15天预测，时间分辨率为6小时，空间分辨率为0.25°，比Gencast快8倍，比IFS-ENS快60倍。它的概率集合表现出始终接近于1的离散技巧比，表明预测的不确定性与观测到的大气变化紧密吻合，从而实现了良好校准的预测。

等级直方图和其他诊断结果证实，集合成员仍然可以与真实世界的观测结果互换，从而证实了FCN3预测的可靠性和可信度。至关重要的是，FCN3保留了所有尺度上的大气频谱特征，忠实地再现了能量级串联和真实世界天气模式的清晰度，即使在长达60天的预测中也是如此。与许多模糊高频特征或随着时间的推移演变成嘈杂伪影的机器学习模型不同，FCN3保持了稳定、物理上真实的频谱，从而能够在亚季节范围内实现准确、清晰且物理上一致的预测。图3显示了FCN3对2020年2月11日初始化的500 hPa风强度的预测，就在丹尼斯风暴在欧洲登陆之前不久。FCN3准确地捕捉了风强度的量级及其在不同长度尺度上的变化，这由各自预测的忠实角度功率谱密度来说明。即使在30天（720小时）或更长的扩展预测中，情况仍然如此。
图片说明

图4：FourCastNet3预测的Atorm Dennis，初始化于2020-02-11 00:00:00 UTC。该图描绘了850hPa压力水平的风速和500hPa地势高度的等值线（高度轮廓）。

FCN3上手指南

经过完整训练的FourCastNet3检查点可在英伟达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性能，我们建议安装启用自定义CUDA扩展的torch-harmonics，并在推理期间使用bf16格式的自动混合精度（这是Earth2Studio中的默认设置）。如果您想运行自定义FCN3推理或自己训练它，您可以在makani中找到代码。
图片说明