Skip to content

Pangu-Weather

None

None

None

# 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. Background Introduction

Pangu-Weather is the first AI method whose accuracy exceeds that of traditional numerical forecasting methods. It provides pre-trained models with 1-hour interval, 3-hour interval, 6-hour interval, and 24-hour interval. The data used includes five meteorological elements (temperature, humidity, geopotential, longitude and latitude wind speeds) on 13 different pressure layers in vertical height, and four meteorological elements on the earth's surface (2-meter temperature, 10-meter wind speed in longitude and latitude directions, sea level pressure). The prediction accuracy from 1 hour to 7 days is higher than that of traditional numerical methods (i.e., operational IFS of the European Meteorological Centre).

At the same time, the Pangu-Weather model can complete a 24-hour global weather forecast in just 1.4 seconds on a V100 graphics card, which is more than 10,000 times faster than traditional numerical forecasting.

2. Model Principle

This chapter only briefly introduces the principle of the Pangu-Weather model. For detailed theoretical derivation, please read Pangu-Weather: A 3D High-Resolution System for Fast and Accurate Global Weather Forecast.

The overall structure of the model is shown in the figure:

result

Model Structure

Its main idea is to use a 3D variant of a visual transformer to process complex and uneven meteorological elements. Due to the high resolution of meteorological data, compared with common vision transformer methods, researchers reduced the encoder and decoder of the network to 2 levels (8 blocks), and adopted the sliding window attention mechanism of Swin transformer to reduce the computation of the network.

The model uses pre-trained weights for inference. Next, the inference process of the model will be introduced.

3. Model Construction

In this case, PanguWeatherPredictor is implemented for inference of the ONNX model:

examples/pangu_weather/predict.py
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
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
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
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'

Among them, input_file and input_surface_file represent the upper-air meteorological data and surface meteorological data input to the network model respectively.

4. Result Visualization

First convert the data from npy to NetCDF format, then use ncvue for visualization

  1. Install dependencies 

    pip install cdsapi netCDF4 ncvue

  2. Use script for data conversion 

    python convert_data.py

  3. Use ncvue to open the converted NetCDF file. For detailed instructions on ncvue, see ncvue official documentation

5. Complete Code

examples/pangu_weather/predict.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
131
132
133
134
135
136
137
138
139
140
141
142
143
# 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. Result Display

The figure below shows the temperature prediction results of the model. More indicators can be viewed using ncvue.

result

Temperature prediction result

7. References