【PPSCI Export&Infer No.34】 (#1129)

mayiru1021 · web-flow · commit aab73e3f047d · 2025-04-09T16:01:54.000+08:00
* pretrain.py添加export和inference

* pretrain.py添加export与infere

* 上次pr修改，风速微调模型及降水量模型添加export与inference，markdown文件添加命令

* precip.yaml文件mode修正

* pretrain.yamlmode修正

* fourcastnet.md文件修正

* 修正0408的review
diff --git a/docs/zh/examples/fourcastnet.md b/docs/zh/examples/fourcastnet.md
@@ -26,6 +26,34 @@
     python train_precip.py mode=eval EVAL.pretrained_model_path=https://paddle-org.bj.bcebos.com/paddlescience/models/fourcastnet/precip.pdparams WIND_MODEL_PATH=https://paddle-org.bj.bcebos.com/paddlescience/models/fourcastnet/finetune.pdparams
     ```
 
+=== "模型导出命令"
+
+    ``` sh
+    # 风速预训练模型导出
+    python train_pretrain.py mode=export
+    # 风速微调模型导出
+    python train_finetune.py mode=export
+    # 降水量模型导出
+    python train_precip.py mode=export
+    ```
+
+=== "模型推理命令"
+
+    ``` sh
+    # 下载风速预测小样本数据
+    wget -nc https://paddle-org.bj.bcebos.com/paddlescience%2Fdatasets%2FFourcastNet%2F2018-04-04_n6_precip.npy -P ./datasets/era5/test/
+    wget -nc https://paddle-org.bj.bcebos.com/paddlescience%2Fdatasets%2FFourcastNet%2F2018-04-04_n6.npy -P ./datasets/era5/test/
+    # 风速预训练模型推理
+    python train_pretrain.py mode=infer
+    # 风速微调模型推理
+    python train_finetune.py mode=infer
+    # 下载降水量预测小样本数据
+    wget -nc https://paddle-org.bj.bcebos.com/paddlescience%2Fdatasets%2FFourcastNet%2F2018-09-08_n32.npy -P ./datasets/era5/test/
+    # 降水量模型推理
+    python train_precip.py mode=infer
+    ```
+
+
 | 模型 | 变量名称 | ACC/RMSE(6h) | ACC/RMSE(30h) | ACC/RMSE(60h) | ACC/RMSE(120h) | ACC/RMSE(192h) |
 | :-- | :-- | :-- | :-- | :-- | :-- | :-- |
 | [风速模型](https://paddle-org.bj.bcebos.com/paddlescience/models/fourcastnet/finetune.pdparams) | U10 | 0.991/0.567 | 0.963/1.130 | 0.891/1.930 | 0.645/3.438 | 0.371/4.915 |
diff --git a/examples/fourcastnet/conf/fourcastnet_finetune.yaml b/examples/fourcastnet/conf/fourcastnet_finetune.yaml
@@ -21,7 +21,7 @@ hydra:
     subdir: ./
 
 # general settings
-mode: train # running mode: train/eval
+mode: train # running mode: train/eval/export/infer
 seed: 1024
 output_dir: ${hydra:run.dir}
 log_freq: 20
@@ -48,10 +48,13 @@ VALID_FILE_PATH: ./datasets/era5/test
 # set test data path
 TEST_FILE_PATH: ./datasets/era5/out_of_sample/2018.h5
 
+INFER_FILE_PATH: ./datasets/era5/test/2018-09-08_n32.npy
+
 # model settings
 MODEL:
   afno:
     input_keys: ["input"]
+    output_keys: ["output"]
 
 # training settings
 TRAIN:
@@ -75,3 +78,22 @@ EVAL:
   compute_metric_by_batch: true
   eval_with_no_grad: true
   batch_size: 1
+
+INFER:
+  pretrained_model_path: "https://paddle-org.bj.bcebos.com/paddlescience/models/fourcastnet/finetune.pdparams"
+  export_path: ./inference/fourcastnet_finetune
+  pdmodel_path: ${INFER.export_path}.pdmodel
+  pdiparams_path: ${INFER.export_path}.pdiparams
+  onnx_path: ${INFER.export_path}.onnx
+  batch_size: ${TRAIN.batch_size}
+  output_keys: ${MODEL.afno.output_keys}
+  device: gpu
+  engine: native
+  precision: fp32
+  ir_optim: true
+  min_subgraph_size: 5
+  gpu_mem: 2000
+  gpu_id: 0
+  max_batch_size: 1024
+  num_cpu_threads: 10
+  num_timestamps: 32
diff --git a/examples/fourcastnet/conf/fourcastnet_precip.yaml b/examples/fourcastnet/conf/fourcastnet_precip.yaml
@@ -21,7 +21,7 @@ hydra:
     subdir: ./
 
 # general settings
-mode: train # running mode: train/eval
+mode: train # running mode: train/eval/export/infer
 seed: 1024
 output_dir: ${hydra:run.dir}
 log_freq: 20
@@ -56,6 +56,10 @@ TEST_FILE_PATH: ./datasets/era5/precip/out_of_sample/2018.h5
 # set wind model path
 WIND_MODEL_PATH: outputs_fourcastnet_finetune/checkpoints/latest
 
+# set inference data path
+WIND_INFER_PATH: ./datasets/era5/test/2018-04-04_n6.npy
+INFER_FILE_PATH: ./datasets/era5/test/2018-04-04_n6_precip.npy
+
 # model settings
 MODEL:
   afno:
@@ -86,3 +90,23 @@ EVAL:
   compute_metric_by_batch: true
   eval_with_no_grad: true
   batch_size: 1
+
+INFER:
+  pretrained_model_path: "https://paddle-org.bj.bcebos.com/paddlescience/models/fourcastnet/precip.pdparams"
+  WIND_MODEL_PATH: "https://paddle-org.bj.bcebos.com/paddlescience/models/fourcastnet/finetune.pdparams"
+  export_path: ./inference/fourcastnet_precip
+  pdmodel_path: ${INFER.export_path}.pdmodel
+  pdiparams_path: ${INFER.export_path}.pdiparams
+  onnx_path: ${INFER.export_path}.onnx
+  batch_size: ${TRAIN.batch_size}
+  output_keys: ${MODEL.afno.output_keys}
+  device: gpu
+  engine: native
+  precision: fp32
+  ir_optim: true
+  min_subgraph_size: 5
+  gpu_mem: 2000
+  gpu_id: 0
+  max_batch_size: 1024
+  num_cpu_threads: 10
+  num_timestamps: 6
diff --git a/examples/fourcastnet/conf/fourcastnet_pretrain.yaml b/examples/fourcastnet/conf/fourcastnet_pretrain.yaml
@@ -21,7 +21,7 @@ hydra:
     subdir: ./
 
 # general settings
-mode: train # running mode: train/eval
+mode: train # running mode: train/eval/export/infer
 seed: 1024
 output_dir: ${hydra:run.dir}
 log_freq: 20
@@ -46,6 +46,9 @@ DATA_TIME_MEAN_PATH: ./datasets/era5/stat/time_means.npy
 # set evaluate data path
 VALID_FILE_PATH: ./datasets/era5/test
 
+# set inference data path
+INFER_FILE_PATH: ./datasets/era5/test/2018-09-08_n32.npy
+
 # model settings
 MODEL:
   afno:
@@ -72,3 +75,23 @@ EVAL:
   compute_metric_by_batch: true
   eval_with_no_grad: true
   batch_size: 8
+
+
+INFER:
+  pretrained_model_path: "https://paddle-org.bj.bcebos.com/paddlescience/models/fourcastnet/pretrain.pdparams"
+  export_path: ./inference/fourcastnet_pretrain
+  pdmodel_path: ${INFER.export_path}.pdmodel
+  pdiparams_path: ${INFER.export_path}.pdiparams
+  onnx_path: ${INFER.export_path}.onnx
+  batch_size: ${TRAIN.batch_size}
+  output_keys: ${MODEL.afno.output_keys}
+  device: gpu
+  engine: native
+  precision: fp32
+  ir_optim: true
+  min_subgraph_size: 5
+  gpu_mem: 2000
+  gpu_id: 0
+  max_batch_size: 1024
+  num_cpu_threads: 10
+  num_timestamps: 32
diff --git a/examples/fourcastnet/train_finetune.py b/examples/fourcastnet/train_finetune.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import functools
+import os
 from os import path as osp
 from typing import Tuple
 
@@ -333,6 +334,110 @@ def output_wind_func(d, var_name, data_mean, data_std):
     solver.visualize()
 
 
+def export(cfg: DictConfig):
+    # set model
+    model = ppsci.arch.AFNONet(**cfg.MODEL.afno)
+
+    # initialize solver
+    solver = ppsci.solver.Solver(
+        model,
+        pretrained_model_path=cfg.INFER.pretrained_model_path,
+    )
+    # export model
+    from paddle.static import InputSpec
+
+    input_spec = [
+        {
+            key: InputSpec([None, 20, cfg.IMG_H, cfg.IMG_W], "float32", name=key)
+            for key in model.input_keys
+        },
+    ]
+    solver.export(input_spec, cfg.INFER.export_path)
+
+
+def inference(cfg: DictConfig):
+    from deploy.python_infer import pinn_predictor
+
+    predictor = pinn_predictor.PINNPredictor(cfg)
+
+    data_mean, data_std = fourcast_utils.get_mean_std(
+        cfg.DATA_MEAN_PATH, cfg.DATA_STD_PATH, cfg.VARS_CHANNEL
+    )
+
+    data = np.load(cfg.INFER_FILE_PATH)
+    input_0 = (data[:, 0] - data_mean) / data_std
+    all_data = input_0
+
+    for t in range(cfg.INFER.num_timestamps):
+        data_t = data[:, t + 1]
+        data_t = (data_t - data_mean) / data_std
+        all_data = np.concatenate((all_data, data_t), axis=0)
+
+    input_dict = {cfg.MODEL.afno.input_keys[0]: all_data}
+
+    vis_output = predictor.predict(input_dict, cfg.INFER.batch_size)
+
+    vis_dict = {
+        store_key: vis_output[infer_key]
+        for store_key, infer_key in zip(cfg.MODEL.afno.output_keys, vis_output.keys())
+    }
+
+    def output_wind_func(output, data_mean, data_std):
+        output = (output * data_std) + data_mean
+        wind_data = (output[0] ** 2 + output[1] ** 2) ** 0.5
+        return wind_data
+
+    wind_pred = []
+    pred_dict = {}
+    for i in range(cfg.INFER.num_timestamps):
+        hour = (i + 1) * 6
+        wind_ = [
+            output_wind_func(
+                vis_dict[cfg.MODEL.afno.output_keys[0]][i], data_mean, data_std
+            )
+        ]
+        wind_pred.append(wind_)
+        pred_dict[f"output_{hour}h"] = np.asarray(wind_)
+    output_dict = {cfg.MODEL.afno.output_keys[0]: np.array(wind_pred)}
+
+    wind_pred = []
+    target_dict = {}
+    for i in range(cfg.INFER.num_timestamps):
+        hour = (i + 1) * 6
+        wind_ = [(data[0][i][0] ** 2 + data[0][i][1] ** 2) ** 0.5]
+        target_dict[f"target_{hour}h"] = np.asarray(wind_)
+
+    vis_dict = {**pred_dict, **target_dict}
+
+    plot_expr_dict = {}
+    for hour in range(6, 6 + cfg.INFER.num_timestamps * 6, 6):
+        plot_expr_dict.update(
+            {
+                f"target_{hour}h": lambda d, hour=hour: d[f"target_{hour}h"],
+                f"output_{hour}h": lambda d, hour=hour: d[f"output_{hour}h"],
+            }
+        )
+
+    visualizer_weather = ppsci.visualize.VisualizerWeather(
+        vis_dict,
+        plot_expr_dict,
+        xticks=np.linspace(0, cfg.IMG_W - 1, 13),
+        xticklabels=[str(i) for i in range(360, -1, -30)],
+        yticks=np.linspace(0, cfg.IMG_H - 1, 7),
+        yticklabels=[str(i) for i in range(90, -91, -30)],
+        vmin=0,
+        vmax=25,
+        colorbar_label="m\s",
+        batch_size=1,
+        num_timestamps=cfg.INFER.num_timestamps,
+        prefix="wind",
+    )
+    visualizer_weather.save(cfg.INFER.export_path, vis_dict)
+    save_path = osp.join(cfg.INFER.export_path, "predict.npy")
+    os.makedirs(cfg.INFER.export_path, exist_ok=True)
+    np.save(save_path, output_dict[cfg.MODEL.afno.output_keys[0]])
+
+
 @hydra.main(
     version_base=None, config_path="./conf", config_name="fourcastnet_finetune.yaml"
 )
@@ -341,8 +446,14 @@ def main(cfg: DictConfig):
         train(cfg)
     elif cfg.mode == "eval":
         evaluate(cfg)
+    elif cfg.mode == "export":
+        export(cfg)
+    elif cfg.mode == "infer":
+        inference(cfg)
     else:
-        raise ValueError(f"cfg.mode should in ['train', 'eval'], but got '{cfg.mode}'")
+        raise ValueError(
+            f"cfg.mode should in ['train', 'eval', 'export', 'infer'], but got '{cfg.mode}'"
+        )
 
 
 if __name__ == "__main__":
diff --git a/examples/fourcastnet/train_precip.py b/examples/fourcastnet/train_precip.py
diff --git a/examples/fourcastnet/train_pretrain.py b/examples/fourcastnet/train_pretrain.py
