跳转至

GraphCast

# linux
wget -c https://paddle-org.bj.bcebos.com/paddlescience/datasets/graphcast/dataset.zip
wget -c https://paddle-org.bj.bcebos.com/paddlescience/datasets/graphcast/dataset-step12.zip
wget -c https://paddle-org.bj.bcebos.com/paddlescience/models/graphcast/params.zip
wget -c https://paddle-org.bj.bcebos.com/paddlescience/models/graphcast/template_graph.zip
wget -c https://paddle-org.bj.bcebos.com/paddlescience/datasets/graphcast/stats.zip
wget -c https://paddle-org.bj.bcebos.com/paddlescience/datasets/graphcast/graphcast-jax2paddle.csv -P ./data/

# curl https://paddle-org.bj.bcebos.com/paddlescience/datasets/graphcast/dataset.zip -o dataset.zip
# curl https://paddle-org.bj.bcebos.com/paddlescience/datasets/graphcast/dataset-step12.zip -o dataset-step12.zip
# curl https://paddle-org.bj.bcebos.com/paddlescience/models/graphcast/template_graph.zip -o template_graph.zip
# curl https://paddle-org.bj.bcebos.com/paddlescience/datasets/graphcast/stats.zip -o stats.zip
# curl https://paddle-org.bj.bcebos.com/paddlescience/datasets/graphcast/graphcast-jax2paddle.csv --create-dirs -o ./data/graphcast-jax2paddle.csv

unzip -q dataset.zip -d data/
unzip -q dataset-step12.zip -d data/
unzip -q params.zip -d data/
unzip -q stats.zip -d data/
unzip -q template_graph.zip -d data/

python graphcast.py mode=eval EVAL.pretrained_model_path="data/params/GraphCast_small---ERA5-1979-2015---resolution-1.0---pressure-levels-13---mesh-2to5---precipitation-input-and-output.pdparams"

1. 背景简介

全球中期天气预报往往是社会和经济领域相关决策的重要依据。传统的数值天气预报模型一般需要通过增加计算资源来提高预报的精度，而无法直接利用历史天气数据来提升基础模型的预测精度。基于机器学习的天气预报模型能够直接利用历史数据训练模型，提升精度，优化计算资源。同时，这种数据驱动的方法使得模型可从数据中的学习到那些不易用显式方程表示的数量关系，从而提高预测的准确性。

GraphCast 是一种基于机器学习的天气预报模型，该模型可以直接从再分析数据中进行训练，并且能够在一分钟内以 0.25° 的分辨率在全球范围内预测超过 10 天的数百个天气变量。论文表明，GraphCast 在 1380 个验证目标的实验中，有 90% 的预测结果显著优于最准确的操作确定性系统（operational deterministic systems），并且模型能很好地预测严重天气事件，包括热带气旋、大气河流和极端温度。

2. 模型原理

\(X^t\) 表示 t 时刻的天气状态预测，

\[ X^{t+1}=GraphCast(X^{t}, X^{t-1}) \]

GraphCast 通过自回归迭代，产生任意长度 T 的预测序列。

\[ X^{t+1:t+T}=(GraphCast(X^{t}, X^{t-1}), GraphCast(X^{t+1}, X^{t}), ... , GraphCast(X^{t+T-1}, X^{t+T-2}))\]

2.1 模型结构

GraphCast 的核心架构采用基于图神经网络（GNN）的“编码‑处理‑解码”结构。基于 GNN 的学习模拟器在学习流体和其他材料的复杂物理动力学方面非常有效，因为它们的表示和计算结构类似于学习型有限元求解器。

GraphCast 的结构图

由于经纬度网格密度是不均匀的，GraphCast 内部不使用经纬度网格，而是使用了“multi-mesh”表示。“multi-mesh”是通过将正二十面体进行 6 次迭代细化来构建的，如下图所示，每次迭代将多面体上的三角面分成 4 个更小的面。

GraphCast 模型运行在图 \(\mathcal{G(V^\mathrm{G}, V^\mathrm{M}, E^\mathrm{M}, E^\mathrm{G2M}, E^\mathrm{M2G})}\) 上。

\(\mathcal{V^\mathrm{G}}\) 是网格点的集合，每个网格节点代表对应经纬度点的大气垂直切片，节点 \(v_𝑖^\mathrm{G}\) 的特征用 \(\mathbf{v}_𝑖^\mathrm{G,features}\) 表示。

\(V^\mathrm{M}\) 是 mesh 节点的集合，mesh 节点是通过将正二十面体迭代划分生成的，节点 \(v_𝑖^\mathrm{M}\) 的特征用 \(\mathbf{v}_𝑖^\mathrm{M,features}\) 表示。

\(\mathcal{E^\mathrm{M}}\) 是一个无向边集合，其中的每条边连接一个发送mesh节点和接收mesh节点，用 \(e^\mathrm{M}_{v^\mathrm{M}_s \rightarrow v^\mathrm{M}_r}\) 表示，对应的特征用 \(\mathbf{e}^\mathrm{M,features}_{v^\mathrm{M}_s \rightarrow v^\mathrm{M}_r}\) 表示。

\(\mathcal{E^\mathrm{G2M}}\) 是一个无向边集合，其中的每条边连接一个发送网格节点和一个接收 mesh 节点，用 \(e^\mathrm{G2M}_{v^\mathrm{G}_s \rightarrow v^M_r}\) 表示，对应的特征用 \(\mathbf{e}^\mathrm{G2M,features}_{v^\mathrm{G}_s \rightarrow v^\mathrm{M}_r}\) 表示。

\(\mathcal{E^\mathrm{M2G}}\) 是一个无向边集合，其中的每条边连接一个发送mesh节点和一个接收网格节点，用 \(e^\mathrm{M2G}_{v^M_s \rightarrow v^G_r}\) 表示，对应的特征用 \(\mathbf{e}^\mathrm{M2G,features}_{v^\mathrm{M}_s \rightarrow v^\mathrm{G}_r}\) 表示。

2.2 编码器

编码器的作用是将数据转化为 GraphCast 内部的数据表示。首先利用五个多层感知机（MLP）将上述五个集合的特征嵌入至内部空间。

\[ \begin{aligned} \mathbf{v}^\mathrm{G}_i = \mathbf{MLP}^\mathrm{embedder}_\mathcal{V^\mathrm{G}}(\mathbf{v}^\mathrm{G,features}_i) \\ \mathbf{v}^\mathrm{M}_i = \mathbf{MLP}^\mathrm{embedder}_\mathcal{V^\mathrm{M}}(\mathbf{v}^\mathrm{M,features}_i) \\ \mathbf{e}^\mathrm{M}_{v^\mathrm{M}_s \rightarrow v^\mathrm{M}_r} = \mathbf{MLP}^\mathrm{embedder}_\mathcal{E^\mathrm{M}}(\mathbf{e}^{\mathrm{M,features}}_{v^\mathrm{M}_s \rightarrow v^\mathrm{M}_r}) \\ \mathbf{e}^\mathrm{G2M}_{v^\mathrm{G}_s \rightarrow v^\mathrm{M}_r} = \mathbf{MLP}^\mathrm{embedder}_\mathcal{E^\mathrm{G2M}}(\mathbf{e}^{\mathrm{G2M,features}}_{v^\mathrm{G}_s \rightarrow v^\mathrm{M}_r}) \\ \mathbf{e}^\mathrm{M2G}_{v^\mathrm{M}_s \rightarrow v^\mathrm{G}_r} = \mathbf{MLP}^\mathrm{embedder}_\mathcal{E^\mathrm{M2G}}(\mathbf{e}^{\mathrm{M2G,features}}_{v^\mathrm{M}_s \rightarrow v^\mathrm{G}_r}) \\ \end{aligned} \]

之后通过一个 Grid2Mesh GNN 层，将信息从网格节点传递到 mesh 节点。\(\mathcal{E^\mathrm{G2M}}\) 中的边通过关联的节点更新信息。

\[ \mathbf{e}^\mathrm{G2M}_{v^\mathrm{G}_s \rightarrow v^\mathrm{M}_r} {'} = \mathbf{MLP}^\mathrm{Grid2Mesh}_\mathcal{E^\mathrm{G2M}}([\mathbf{e}^\mathrm{G2M}_{v^\mathrm{G}_s \rightarrow v^\mathrm{M}_r}, \mathbf{v}_r^\mathrm{G}, \mathbf{v}_s^\mathrm{M}]) \]

mesh 节点通过其关联的边更新信息。

\[ \mathbf{v}^\mathrm{M}_i {'} = \mathbf{MLP}^\mathrm{Grid2Mesh}_\mathcal{V^\mathrm{M}}([\mathbf{v}^\mathrm{M}_i, \sum_{\mathbf{e}^\mathrm{G2M}_{v^\mathrm{G}_s \rightarrow v^\mathrm{M}_r} : v^\mathrm{M}_r=v^\mathrm{M}_i} \mathbf{e}^\mathrm{G2M}_{v^\mathrm{G}_s \rightarrow v^\mathrm{M}_r} {'}]) \]

同样网格节点也进行信息更新。

\[ \mathbf{v}^\mathrm{G}_i {'} = \mathbf{MLP}^\mathrm{Grid2Mesh}_\mathcal{V^\mathrm{G}}(\mathbf{v}^\mathrm{G}_i) \]

最后通过残差连接更新三个元素。

\[ \begin{aligned} \mathbf{v}^\mathrm{G}_i \leftarrow \mathbf{v}^\mathrm{G}_i + \mathbf{v}^\mathrm{G}_i {'} \\ \mathbf{v}^\mathrm{M}_i \leftarrow \mathbf{v}^\mathrm{M}_i + \mathbf{v}^\mathrm{M}_i {'} \\ \mathbf{e}^\mathrm{G2M}_{v^\mathrm{G}_s \rightarrow v^\mathrm{M}_r} = \mathbf{e}^\mathrm{G2M}_{v^\mathrm{G}_s \rightarrow v^\mathrm{M}_r} + \mathbf{e}^\mathrm{G2M}_{v^\mathrm{G}_s \rightarrow v^\mathrm{M}_r} {'} \end{aligned} \]

2.3 处理器

处理器包含一个Multi-mesh GNN 层，\(\mathcal{E^\mathrm{M}}\) 中的边通过关联的节点更新信息。

\[ \mathbf{e}^\mathrm{M}_{v^\mathrm{M}_s \rightarrow v^\mathrm{M}_r} {'} = \mathbf{MLP}^\mathrm{Mesh}_\mathcal{E^\mathrm{M}}([\mathbf{e}^\mathrm{M}_{v^\mathrm{M}_s \rightarrow v^\mathrm{M}_r}, \mathbf{v}^\mathrm{M}_s, \mathbf{v}^\mathrm{M}_r]) \]

mesh 节点通过其关联的边更新信息。

\[ \mathbf{v}^\mathrm{M}_i {'} = \mathbf{MLP}^\mathrm{Mesh}_\mathcal{V^\mathrm{M}}([\mathbf{v}^\mathrm{M}_i, \sum_{\mathbf{e}^\mathrm{G2M}_{v^\mathrm{G}_s \rightarrow v^\mathrm{M}_r} : v^\mathrm{M}_r=v^\mathrm{M}_i} \mathbf{e}^\mathrm{M}_{v^\mathrm{G}_s \rightarrow v^\mathrm{M}_r} {'}]) \]

最后通过残差连接更新元素。

\[ \begin{aligned} \mathbf{v}^\mathrm{M}_i \leftarrow \mathbf{v}^\mathrm{M}_i + \mathbf{v}^\mathrm{M}_i {'} \\ \mathbf{e}^\mathrm{M}_{v^\mathrm{M}_s \rightarrow v^\mathrm{M}_r} \leftarrow \mathbf{e}^\mathrm{M}_{v^\mathrm{M}_s \rightarrow v^\mathrm{M}_r} + \mathbf{e}^\mathrm{M}_{v^\mathrm{M}_s \rightarrow v^\mathrm{M}_r} {'}\\ \end{aligned} \]

2.4 解码器

解码器的作用是将 mesh 内的信息取回网格中，并进行预测。解码器包含一个Mesh2Grid GNN 层。

\(\mathcal{E^\mathrm{M2G}}\) 中的边通过关联的节点的更新信息。

\[ \mathbf{e}^\mathrm{M2G}_{v^\mathrm{M}_s \rightarrow v^\mathrm{G}_r} {'} = \mathbf{MLP}^\mathrm{Mesh2Grid}_\mathcal{E^\mathrm{M2G}}([\mathbf{e}^\mathrm{M2G}_{v^\mathrm{M}_s \rightarrow v^\mathrm{G}_r},\mathbf{v}^\mathrm{M}_s, \mathbf{v}^\mathrm{M}_r]) \]

网格节点通过其关联的边更新信息。

\[ \mathbf{v}^\mathrm{G}_i {'} = \mathbf{MLP}^\mathrm{Mesh2Grid}_\mathcal{V^\mathrm{G}}([\mathbf{v}^\mathrm{G}_i, \sum_{\mathbf{e}^\mathrm{G2M}_{v^\mathrm{M}_s \rightarrow v^\mathrm{G}_r} : v^\mathrm{G}_r=v^\mathrm{G}_i} \mathbf{e}^\mathrm{M2G}_{v^\mathrm{M}_s \rightarrow v^\mathrm{G}_r} {'}]) \]

通过残差连接对网格节点进行更新。

\[ \mathbf{v}^\mathrm{G}_i \leftarrow \mathbf{v}^\mathrm{G}_i + \mathbf{v}^\mathrm{G}_i {'} \]

接着利用另一个 MLP 对网格信息进行处理，得到预测值。

\[ \mathbf{\hat{y}}^\mathrm{G}_i= \mathbf{MLP}^\mathrm{Output}_\mathcal{V^\mathrm{G}}(\mathbf{v}^\mathrm{G}_i) \]

将输入状态 \(X^{t}\) 与预测值 \(\hat{Y}^{t}\) 相加得到下一个天气状态 \(\hat{X}^{t+1}\)

\[ \hat{X}^{t+1} = GraphCast(X^{t}, X^{t-1}) = X^{t} + \hat{Y}^{t} \]

3. 模型构建

接下来开始讲解如何基于 PaddleScience 代码，实现 GraphCast。关于该案例中的其余细节请参考 API文档。

3.1 数据集介绍

数据集采用了 ECMWF 的 ERA5 数据集 的 2020年再分析存档子集，数据时间段为1979-2018 年，时间间隔为6小时（对应每天的00z、06z、12z和18z），水平分辨率为0.25°，包含 37 个垂直大气压力层。

模型预测总共227个目标变量，其中包括5个地面变量，以及在13个压力层中的每个层次的6个大气变量。

3.2 加载预训练模型

在执行命令中设定预训练模型的文件路径，如下。

python graphcast.py mode=eval EVAL.pretrained_model_path="data/params/GraphCast_small---ERA5-1979-2015---resolution-1.0---pressure-levels-13---mesh-2to5---precipitation-input-and-output.pdparams"

3.3 模型构建

我们使用神经网络 GraphCastNet 作为模型，其接收天气数据，输出预测结果。

28
model = ppsci.arch.GraphCastNet(**cfg.MODEL)

3.4 评估器构建

我们使用 ppsci.validate.SupervisedValidator 构建评估器。首先定义数据加载器的配置，然后创建评估器。

31
32
33
34
35
36
37
38
39
eval_dataloader_cfg = {
    "dataset": {
        "name": "GridMeshAtmosphericDataset",
        "input_keys": ("input",),
        "label_keys": ("label",),
        **cfg.DATA,
    },
    "batch_size": cfg.EVAL.batch_size,
}

我们需要定义训练损失函数的计算过程。

50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
def loss(
    output_dict: Dict[str, paddle.Tensor],
    label_dict: Dict[str, paddle.Tensor],
    *args,
) -> Dict[str, paddle.Tensor]:
    graph = output_dict["pred"]
    pred = dataset.denormalize(graph.grid_node_feat.numpy())
    pred = graph.grid_node_outputs_to_prediction(pred, dataset.targets_template)

    target = graph.grid_node_outputs_to_prediction(
        label_dict["label"][0].numpy(), dataset.targets_template
    )

    pred = atmospheric_dataset.dataset_to_stacked(pred)
    target = atmospheric_dataset.dataset_to_stacked(target)
    loss = np.average(np.square(pred.data - target.data))
    loss = paddle.full([], loss)
    return {"loss": loss}

接着我们还需要定义 metric 指标。

69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
def metric(
    output_dict: Dict[str, paddle.Tensor],
    label_dict: Dict[str, paddle.Tensor],
    *args,
) -> Dict[str, paddle.Tensor]:
    graph = output_dict["pred"][0]
    pred = dataset.denormalize(graph.grid_node_feat.numpy())
    pred = graph.grid_node_outputs_to_prediction(pred, dataset.targets_template)

    target = graph.grid_node_outputs_to_prediction(
        label_dict["label"][0].numpy(), dataset.targets_template
    )

    metric_dic = {
        var_name: np.average(target[var_name].data - pred[var_name].data)
        for var_name in list(target)
    }
    return metric_dic

最后完成评估器的构建。

88
89
90
91
92
dataset = error_validator.data_loader.dataset
error_validator.loss = ppsci.loss.FunctionalLoss(loss)
error_validator.metric = {"error": ppsci.metric.FunctionalMetric(metric)}

validator = {error_validator.name: error_validator}

3.5 模型评估

完成上述设置之后，只需要将上述实例化的对象按顺序传递给 ppsci.solver.Solver，然后启动评估。

 94
 95
 96
 97
 98
 99
100
101
102
103
104
# initialize solver
solver = ppsci.solver.Solver(
    model,
    validator=validator,
    cfg=cfg,
    pretrained_model_path=cfg.EVAL.pretrained_model_path,
    eval_with_no_grad=cfg.EVAL.eval_with_no_grad,
)

# evaluate model
solver.eval()

3.6 结果可视化

评估完成后，我们以图片的形式对结果进行可视化，如下所示。

106
107
108
109
110
111
112
113
114
115
116
117
118
# visualize prediction
with solver.no_grad_context_manager(True):
    for index, (input_, label_, _) in enumerate(error_validator.data_loader):
        output_ = model(input_)
        graph = output_["pred"]
        pred = dataset.denormalize(graph.grid_node_feat.numpy())
        pred = graph.grid_node_outputs_to_prediction(pred, dataset.targets_template)

        target = graph.grid_node_outputs_to_prediction(
            label_["label"][0].numpy(), dataset.targets_template
        )

        plot.log_images(target, pred, "2m_temperature", level=50, file="result.png")

4. 完整代码

graphcast.py
  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

#     http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import Dict

import hydra
import numpy as np
import paddle
import plot
from omegaconf import DictConfig

import ppsci
from ppsci.data.dataset import atmospheric_dataset


def eval(cfg: DictConfig):
    model = ppsci.arch.GraphCastNet(**cfg.MODEL)

    # set dataloader config
    eval_dataloader_cfg = {
        "dataset": {
            "name": "GridMeshAtmosphericDataset",
            "input_keys": ("input",),
            "label_keys": ("label",),
            **cfg.DATA,
        },
        "batch_size": cfg.EVAL.batch_size,
    }

    # set validator
    error_validator = ppsci.validate.SupervisedValidator(
        eval_dataloader_cfg,
        loss=None,
        output_expr={"pred": lambda out: out["pred"]},
        metric=None,
        name="error_validator",
    )

    def loss(
        output_dict: Dict[str, paddle.Tensor],
        label_dict: Dict[str, paddle.Tensor],
        *args,
    ) -> Dict[str, paddle.Tensor]:
        graph = output_dict["pred"]
        pred = dataset.denormalize(graph.grid_node_feat.numpy())
        pred = graph.grid_node_outputs_to_prediction(pred, dataset.targets_template)

        target = graph.grid_node_outputs_to_prediction(
            label_dict["label"][0].numpy(), dataset.targets_template
        )

        pred = atmospheric_dataset.dataset_to_stacked(pred)
        target = atmospheric_dataset.dataset_to_stacked(target)
        loss = np.average(np.square(pred.data - target.data))
        loss = paddle.full([], loss)
        return {"loss": loss}

    def metric(
        output_dict: Dict[str, paddle.Tensor],
        label_dict: Dict[str, paddle.Tensor],
        *args,
    ) -> Dict[str, paddle.Tensor]:
        graph = output_dict["pred"][0]
        pred = dataset.denormalize(graph.grid_node_feat.numpy())
        pred = graph.grid_node_outputs_to_prediction(pred, dataset.targets_template)

        target = graph.grid_node_outputs_to_prediction(
            label_dict["label"][0].numpy(), dataset.targets_template
        )

        metric_dic = {
            var_name: np.average(target[var_name].data - pred[var_name].data)
            for var_name in list(target)
        }
        return metric_dic

    dataset = error_validator.data_loader.dataset
    error_validator.loss = ppsci.loss.FunctionalLoss(loss)
    error_validator.metric = {"error": ppsci.metric.FunctionalMetric(metric)}

    validator = {error_validator.name: error_validator}

    # initialize solver
    solver = ppsci.solver.Solver(
        model,
        validator=validator,
        cfg=cfg,
        pretrained_model_path=cfg.EVAL.pretrained_model_path,
        eval_with_no_grad=cfg.EVAL.eval_with_no_grad,
    )

    # evaluate model
    solver.eval()

    # visualize prediction
    with solver.no_grad_context_manager(True):
        for index, (input_, label_, _) in enumerate(error_validator.data_loader):
            output_ = model(input_)
            graph = output_["pred"]
            pred = dataset.denormalize(graph.grid_node_feat.numpy())
            pred = graph.grid_node_outputs_to_prediction(pred, dataset.targets_template)

            target = graph.grid_node_outputs_to_prediction(
                label_["label"][0].numpy(), dataset.targets_template
            )

            plot.log_images(target, pred, "2m_temperature", level=50, file="result.png")


@hydra.main(version_base=None, config_path="./conf", config_name="graphcast_small.yaml")
def main(cfg: DictConfig):
    if cfg.mode == "eval":
        eval(cfg)
    else:
        raise ValueError(f"cfg.mode should in ['eval'], but got '{cfg.mode}'")


if __name__ == "__main__":
    main()

5. 结果展示

下图展示了温度的真值结果、预测结果和误差。

result_wind

真值结果（"targets"）、预测结果（"prediction"）和误差（"diff"）

可以看到模型预测结果与真实结果基本一致。

6. 参考文献