GraphCast天气预测：图神经网络GNN天气

可视化演示

展示球面网格图构建与消息传递机制、多变量多气压层大气特征编码、数据并行与模型并行训练策略，交互切换核心环节。

图结构构建：画一张全村地图

要组织全村人互通天气，第一件事是画地图：把村庄分成一个个小院，每院标好位置，用道路把相邻的院落连起来。这张“院落为点、道路为线”的地图，就是图神经网络里说的“图结构”。没有地图，谁和谁相邻都说不清，消息传递更是无从谈起。

对应到 GraphCast，地图上的“院落”来自对地球表面的网格剖分。模型以 0.25° 经度×纬度为一个单元，把球面划分成 721×1440 个网格，约 103.8 万个节点；相邻网格之间连一条无向边，表示“这两个院落地理会相互影响”。每个节点用经纬度定位，就像一个院子在村庄中的坐标。

但地球上的天气系统尺度很大——台风的外围水汽可能来自数百公里外。光有细网格地图不够，GraphCast 于是又叠加了两档更粗的地图：1° 分辨率和 2.25° 分辨率。细地图上的短边照顾近邻，粗地图上的长边把信息送到更远，三层地图合在一起构成多尺度图。数学上，整张图可以写成：

有了这幅多层次的地图，信息就有了流动的通道。但下一个问题随之而来：每个村民的纸条上到底该写什么？这正是特征编码要回答的事。

特征编码方法：村民在纸条上写什么

地图画好了，村民拿起笔，把自家观测整理成一张标准纸条。纸条上的内容越丰富，邻居们做判断时依据就越充分。但也不能什么都往上写——写得太多、格式不统一，反而让人抓不住重点。

GraphCast 的“纸条”是一串气象变量向量：37 个气压层上的温度、相对湿度、东西向风速、南北向风速、重力位势，加上地表 2 米气温、10 米风速、海平面气压等，合计 189 个大气状态变量。此外还有海拔、海陆掩膜、纬度等“静态特征”——它们几乎不随时间变化，却决定了这个院落的先天属性，比如山区还是沿海、低海拔还是高海拔。

问题来了：温度的数值在十几到几十之间，气压的数值却在 1000 百帕量级。如果村民直接写原始数字，“大数字”会淹没“小数字”，模型很难公平地看待每个变量。所以大家先做归一化——减去历史均值、除以历史标准差：

纸条内容准备好了，村民还得把它翻译成“村里的统一语言”——也就是嵌入表示，这个细节留到编码器设计一节再展开。先看最关键的问题：纸条究竟是怎么在邻居之间传递的？

消息传递机制：纸条如何流动

全村人坐进会议厅，递纸条的规则有三步：先写消息，再汇总，最后更新判断。这套规则就是 GNN 的核心——消息传递。每个节点把自己当前的状态写成一条条消息，沿着边送给邻居；每个节点把所有收到的消息汇总，结合自己原来的状态，更新出新的状态。一轮下来，每个节点就“听见”了所有直接邻居的声音。

在 GraphCast 的每一层网络里，这三步由三个可学习的函数完成。发送方节点 u 写下消息，接收方节点 v 汇总邻居们的消息，再更新自己的状态：

一轮消息传递只让信息在“直接邻居”间流动，但天气的影响可以跨越数千公里。GraphCast 把上述三步连续做 16 层——每一轮，信息就顺着边向外扩散一层。就像村民会议开了一轮又一轮，消息从邻院传到几条街外、再传到几个村子外。掌握了这个核心机制，接下来就能把完整流水线拆开看：纸条怎么被格式化（编码器）、怎么在会议厅里反复讨论（处理器）、最后怎么宣读结论（解码器）。

编码器设计：把纸条翻译成标准格式

会议前，每位村民手里攥着十几条原始观测数据，长短不一、格式混乱。直接这样是没法开会的。所以村里定下规矩：所有纸条必须先翻译成统一格式的“编码”，登记在相同式样的本子上。这就是编码器的作用——把原始观测变成机器能高效处理的嵌入向量。

编码器把节点 v 的原始特征向量 x_v（包括 189 个气象变量、静态特征和时间特征）通过一个可学习的线性映射，变成高维嵌入向量 h_v^{(0)}，作为第一层消息传递的起点：

GraphCast 的编码器还承担“网格到图”的桥梁作用：三个分辨率网格上的每个单元都要逐一映射成图节点，带上自己的嵌入向量；同时把边界条件等额外信息一并融合进来。经过这一步，103 万个网格单元不再是孤立的坐标点，而是一个个拥有统一“语言”的会议参与者，准备进入最深层的处理器。

处理器核心：反复开 16 场村民会议

纸条格式统一了，正式的“社区会议”开始。处理器是整条流水线中最深的部分——16 层 GNN 堆叠，每一层都是一次完整的消息传递。但深层网络有一个经典难题：层数越深，梯度传播越不稳定，模型的效果反而会退化。

村民开会也是如此——如果每轮讨论都推倒重来，16 轮下来肯定一团乱麻。要让每一轮都在上一轮结论上“打补丁”，就需要残差连接：把上一层的状态原样保留，加上本轮 GNN 学到的变化量：

处理器内部还有一个精妙设计：16 层 GNN 共享同一套参数。这就像全村统一一套议事规则——16 轮会议议题不断变化，但大家遵循的讨论逻辑是同一套，参数量因此大幅减少，泛化能力也更强。经过 16 轮讨论，每个节点都收敛到了一个融合了远近信息的隐状态，准备进入最后阶段。

解码器输出：宣读最终预报

会议结束，村民回到自家院落，手里攥着只有自己看得懂的“会议纪要”——高维隐状态向量。但邻居们看不懂这些抽象符号。解码器的任务，就是把这些隐状态翻译回人人都能懂的天气数字：温度几度、风速几米每秒、会不会下雨。

解码器是编码器的镜像：先把图节点的隐状态映射回原始网格坐标（图到网格），再通过一个输出头产生两类结果——下一时刻的完整大气状态（6 小时后的温度、风速、湿度等），以及诊断量（如降水概率）：

一次前向推演只得到 6 小时预报。要拿到 10 天预报，模型采用“迭代推演”：把刚预测出的 6 小时状态当作新输入，再推下一个 6 小时……循环 40 次。这种“滚雪球”式的自回归预测，让每一步都基于模型自己的最新预测。GraphCast 正是靠这个方案，在约 60 秒内完成 10 天全球预报，精度在 90% 以上的验证指标上超过了传统数值预报系统 HRES。

计算图优化：让村庄会议更高效

103 万个节点、上千万条边，每层消息传递都要在边上做运算——如果不加优化，计算量和显存开销都是天文数字。就像村里开会，如果每个人都给全村人写信，邮局会直接瘫痪。聪明的办法是：只给邻居写信，并把所有信件打包成邮袋集中投递。

第一招是边稀疏化。GraphCast 在 0.25° 细网格上只保留距离约 316 公里以内的边，更远的依赖交给中、粗网格上的长边，每个节点平均只有几十条边。第二招是算子融合：把所有边上的消息计算合并成一次大型矩阵乘法，充分利用 GPU 的并行能力，而不是逐条边循环：

GraphCast 基于 JAX 实现，XLA 编译器会对整张计算图做进一步优化：自动融合算子、消除中间缓冲、调整执行顺序。显存管理上也有巧思——训练时把 103 万节点的上采样、下采样放到不同设备上，按需分配显存。经过这些优化，原本“不可能”的大规模 GNN 天气预报才真正落地。

并行训练方法：全村人一起学历史

模型再精巧，也得有“老师”教。GraphCast 的老师是 ECMWF 提供的 ERA5 再分析资料——1979 到 2017 年整整 39 年的全球历史天气场。如何用这批数据教会 103 万节点上的 GNN？答案是：让多个村庄（GPU）同时开工。

最直观的思路是数据并行：把 39 年的样本切成多份，每张 GPU 学一份，各自算出梯度再汇总更新。但图太大时，一张卡存不下完整的图，就需要图切分：把地球表面的节点按空间区块切开分到多张卡上，卡与卡之间只交换切口附近的边界节点状态。

当显存仍然吃紧时，还可以用梯度累积。并不是每个小批次都要立刻更新模型，而是先攒若干个批次的梯度，再统一更新一次参数：

分布式通信同样讲究：卡与卡之间传递的不是整张图，而是边界上的少量节点嵌入，通信量大幅压缩。最终，这套并行训练方案让 GraphCast 在 32 块 TPU v4 上约 4 周完成训练——“学会 39 年全球天气”从不可能变成了现实。

AI 应用场景

Python 演示

import numpy as np  def build_grid_graph(height, width):    """构建 4 邻接网格图，返回边列表"""    edges = []    for r in range(height):        for c in range(width):            idx = r * width + c            if r > 0:                edges.append((idx, idx - width))            if r < height - 1:                edges.append((idx, idx + width))            if c > 0:                edges.append((idx, idx - 1))            if c < width - 1:                edges.append((idx, idx + 1))    return np.array(edges)  def message_passing(features, edges, w_msg, w_update, num_layers=3):    """多层消息传递：邻居消息聚合后更新节点状态"""    h = features.copy()    degree = np.zeros(features.shape[0])    for i, j in edges:        degree[j] += 1.0    for _ in range(num_layers):        messages = np.zeros_like(h)        for i, j in edges:            messages[j] += h[i] @ w_msg        messages /= np.maximum(degree[:, None], 1.0)        h = h @ w_update + messages    return h  def main():    np.random.seed(42)    height, width = 4, 4    x = np.linspace(0, 2 * np.pi, width)    y = np.linspace(0, 2 * np.pi, height)    grid_x, grid_y = np.meshgrid(x, y)    temperature = 15.0 + 5.0 * np.sin(grid_x) + 3.0 * np.cos(grid_y)    land = np.ones((height, width))    land[1:3, 1:3] = 0.0    node_feat = np.stack([temperature / 30.0, land], axis=-1).reshape(-1, 2)     edges = build_grid_graph(height, width)    w_msg = np.random.randn(2, 2).astype(np.float32) * 0.1    w_update = np.random.randn(2, 2).astype(np.float32) * 0.1     forecast = message_passing(node_feat, edges, w_msg, w_update, num_layers=3)    print("初始温度场(摄氏度):")    print(np.round(temperature, 2))    print("6小时后预测温度场(摄氏度):")    print(np.round(forecast[:, 0] * 30.0, 2).reshape(height, width))  if __name__ == "__main__":    main()

1. edges.append((idx, idx - width))为每个节点连上“北边邻居”，构建邻接边列表，相当于给村庄画道路。
2. messages[j] += h[i] @ w_msg邻居 i 的消息乘以权重矩阵后累加到节点 j，即“消息生成 + 邻居聚合”一步完成。
3. messages /= np.maximum(degree[:, None], 1.0)按邻居数量归一化，防止高连接度节点聚合值过大，是图上的“归一化”。
4. h = h @ w_update + messages节点用自身旧状态乘以更新矩阵，再加上邻居汇总消息，完成“节点更新”。
5. node_feat = np.stack([temperature / 30.0, land], axis=-1)把温度场和陆地掩膜拼成节点特征向量，同时做缩放到同数量级，对应特征编码。

常见误区

面试题