Pangu-Weather¶

模型训练命令模型评估命令模型导出命令模型推理命令

暂无

# Download sample input data
wget -c https://paddle-org.bj.bcebos.com/paddlescience/models/Pangu/input_surface.npy -P ./data
wget -c https://paddle-org.bj.bcebos.com/paddlescience/models/Pangu/input_upper.npy -P ./data

# Download pretrain model weight
wget -c https://paddle-org.bj.bcebos.com/paddlescience/models/Pangu/pangu_weather_1.onnx -P ./inference
wget -c https://paddle-org.bj.bcebos.com/paddlescience/models/Pangu/pangu_weather_3.onnx -P ./inference
wget -c https://paddle-org.bj.bcebos.com/paddlescience/models/Pangu/pangu_weather_6.onnx -P ./inference
wget -c https://paddle-org.bj.bcebos.com/paddlescience/models/Pangu/pangu_weather_24.onnx -P ./inference

# 1h interval-time model inference
python predict.py INFER.export_path=inference/pangu_weather_1
# 3h interval-time model inference
python predict.py INFER.export_path=inference/pangu_weather_3
# 6h interval-time model inference
python predict.py INFER.export_path=inference/pangu_weather_6
# 24h interval-time model inference
python predict.py INFER.export_path=inference/pangu_weather_24

1. 背景简介¶

盘古气象大模型(Pangu-Weather)是首个精度超过传统数值预报方法的 AI 方法，其提供了 1 小时间隔、3 小时间隔、6 小时间隔、24 小时间隔的预训练模型。其使用的数据，包括垂直高度上13个不同气压层，每层五种气象要素（温度、湿度、位势、经度和纬度方向的风速），以及地球表面的四种气象要素（2米温度、经度和纬度方向的10米风速、海平面气压）。1 小时 - 7 天预测精度均高于传统数值方法（即欧洲气象中心的 operational IFS）。

同时，盘古气象大模型在一张V100显卡上只需要1.4秒就能完成24小时的全球气象预报，相比传统数值预报提速10000倍以上。

2. 模型原理¶

本章节仅对盘古气象大模型的原理进行简单地介绍，详细的理论推导请阅读 Pangu-Weather: A 3D High-Resolution System for Fast and Accurate Global Weather Forecast。

模型的总体结构如图所示：

其主要思想是使用一个视觉transformer的3D变种来处理复杂的不均匀的气象要素。由于气象数据分辨率很大，因而相比于常见的vision transformer方法，研究人员将网络的encoder和decoder减少到2级（8个block），同时采用Swin transformer的滑窗注意力机制，以减少网络的计算量

模型使用预训练权重推理，接下来将介绍模型的推理过程。

3. 模型构建¶

在该案例中，实现了 PanguWeatherPredictor用于ONNX模型的推理：

examples/pangu_weather/predict.py
def predict(
    self,
    input_data: np.ndarray,
    input_surface_data: np.ndarray,
    batch_size: int = 1,
) -> Tuple[np.ndarray, np.ndarray]:
    """Predicts the output of the yinglong model for the given input.

    Args:
        input_data (np.ndarray): Input data.
        input_surface_data (np.ndarray): Input Surface data.
        batch_size (int, optional): Batch size, now only support 1. Defaults to 1.

    Returns:
        Tuple[np.ndarray, np.ndarray]: Prediction.
    """
    if batch_size != 1:
        raise ValueError(
            f"PanguWeatherPredictor only support batch_size=1, but got {batch_size}"
        )

    # prepare input dict
    input_dict = {
        self.input_names[0]: input_data,
        self.input_names[1]: input_surface_data,
    }

    # run predictor
    output_data, output_surface_data = self.predictor.run(None, input_dict)

    return output_data, output_surface_data

examples/pangu_weather/conf/pangu_weather.yaml
INFER:
  pretrained_model_path: null
  export_path: inference/pangu_weather_24
  onnx_path: ${INFER.export_path}.onnx
  device: gpu
  engine: onnx
  precision: fp32
  ir_optim: false
  min_subgraph_size: 30
  gpu_mem: 100
  gpu_id: 0
  max_batch_size: 1
  num_cpu_threads: 10
  batch_size: 1
  input_file: './data/input_upper.npy'
  input_surface_file: './data/input_surface.npy'

其中，input_file 和 input_surface_file 分别代表网络模型输入的高空气象数据和地面气象。

4. 结果可视化¶

先将数据从 npy 转换为 NetCDF 格式，然后采用 ncvue 进行可视化

安装相关依赖
```
pip install cdsapi netCDF4 ncvue
```
使用脚本进行数据转换
```
python convert_data.py
```
使用 ncvue 打开转换后的 NetCDF 文件, ncvue 具体说明见ncvue官方文档

5. 完整代码¶

examples/pangu_weather/predict.py
# Copyright (c) 2025 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 os import path as osp
from typing import Tuple

import hydra
import numpy as np
import paddle
from omegaconf import DictConfig
from packaging import version

from deploy.python_infer import base
from ppsci.utils import logger


class PanguWeatherPredictor(base.Predictor):
    """General predictor for PanguWeather model.

    Args:
        cfg (DictConfig): Running configuration.
    """

    def __init__(
        self,
        cfg: DictConfig,
    ):
        assert cfg.INFER.engine == "onnx", "Pangu-Weather engine only supports 'onnx'."

        super().__init__(
            pdmodel_path=None,
            pdiparams_path=None,
            device=cfg.INFER.device,
            engine=cfg.INFER.engine,
            precision=cfg.INFER.precision,
            onnx_path=cfg.INFER.onnx_path,
            ir_optim=cfg.INFER.ir_optim,
            min_subgraph_size=cfg.INFER.min_subgraph_size,
            gpu_mem=cfg.INFER.gpu_mem,
            gpu_id=cfg.INFER.gpu_id,
            max_batch_size=cfg.INFER.max_batch_size,
            num_cpu_threads=cfg.INFER.num_cpu_threads,
        )
        self.log_freq = cfg.log_freq

        # get input names
        self.input_names = [
            input_node.name for input_node in self.predictor.get_inputs()
        ]

        # get output names
        self.output_names = [
            output_node.name for output_node in self.predictor.get_outputs()
        ]

    def predict(
        self,
        input_data: np.ndarray,
        input_surface_data: np.ndarray,
        batch_size: int = 1,
    ) -> Tuple[np.ndarray, np.ndarray]:
        """Predicts the output of the yinglong model for the given input.

        Args:
            input_data (np.ndarray): Input data.
            input_surface_data (np.ndarray): Input Surface data.
            batch_size (int, optional): Batch size, now only support 1. Defaults to 1.

        Returns:
            Tuple[np.ndarray, np.ndarray]: Prediction.
        """
        if batch_size != 1:
            raise ValueError(
                f"PanguWeatherPredictor only support batch_size=1, but got {batch_size}"
            )

        # prepare input dict
        input_dict = {
            self.input_names[0]: input_data,
            self.input_names[1]: input_surface_data,
        }

        # run predictor
        output_data, output_surface_data = self.predictor.run(None, input_dict)

        return output_data, output_surface_data


def inference(cfg: DictConfig):
    # log paddlepaddle's version
    if version.Version(paddle.__version__) != version.Version("0.0.0"):
        paddle_version = paddle.__version__
        if version.Version(paddle.__version__) < version.Version("2.6.0"):
            logger.warning(
                f"Detected paddlepaddle version is '{paddle_version}', "
                "currently it is recommended to use release 2.6 or develop version."
            )
    else:
        paddle_version = f"develop({paddle.version.commit[:7]})"

    logger.info(f"Using paddlepaddle {paddle_version}")

    # create predictor
    predictor = PanguWeatherPredictor(cfg)

    # load data
    input_data = np.load(cfg.INFER.input_file).astype(np.float32)
    input_surface_data = np.load(cfg.INFER.input_surface_file).astype(np.float32)

    # run predictor
    output_data, output_surface_data = predictor.predict(input_data, input_surface_data)

    # save predict data
    output_save_path = osp.join(cfg.output_dir, "output_upper.npy")
    np.save(output_save_path, output_data)
    output_surface_save_path = osp.join(cfg.output_dir, "output_surface.npy")
    np.save(output_surface_save_path, output_surface_data)
    logger.info(
        f"Save output upper to {output_save_path} and output surface to {output_surface_save_path}."
    )


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


if __name__ == "__main__":
    main()

6. 结果展示¶

下图展示了模型的温度预测结果，更多指标可以使用 ncvue 查看。

7. 参考资料¶

Pangu-Weather: A 3D High-Resolution System for Fast and Accurate Global Weather Forecast