【PPSCI Export&Infer No.8】Add export and inference for amgnet_cylinder

.gitignore

@@ -132,3 +132,8 @@ FETCH_HEAD
 
 # auto generated version file by setuptools_scm
 ppsci/_version.py
+examples/amgnet/data/
+examples/amgnet/data.zip
+examples/amgnet/inference/
+examples/amgnet/result/
+examples/amgnet/outputs_amgnet_*/

docs/zh/examples/amgnet.md

@@ -56,6 +56,22 @@
         python amgnet_cylinder.py mode=eval EVAL.pretrained_model_path=https://paddle-org.bj.bcebos.com/paddlescience/models/amgnet/amgnet_cylinder_pretrained.pdparams
         ```
 
+=== "模型导出命令"
+
+    === "amgnet_cylinder"
+
+        ``` sh
+        python amgnet_cylinder.py mode=export
+        ```
+
+=== "Python推理命令"
+
+    === "amgnet_cylinder"
+
+        ``` sh
+        python amgnet_cylinder.py mode=infer
+        ```
+
 | 预训练模型  | 指标 |
 |:--| :--|
 | [amgnet_airfoil_pretrained.pdparams](https://paddle-org.bj.bcebos.com/paddlescience/models/amgnet/amgnet_airfoil_pretrained.pdparams) | loss(RMSE_validator): 0.0001 <br> RMSE.RMSE(RMSE_validator): 0.01315 |
@@ -291,6 +307,64 @@ unzip data.zip
     --8<--
     ```
 
+### 3.9 模型导出与推理
+
+训练完成后，我们可以将模型导出为静态图格式，并使用Python推理引擎进行部署。
+
+#### 3.9.1 导出模型
+
+我们首先需要在 `amgnet_cylinder.py` 中实现 `export` 函数，它负责加载训练好的模型，并将其保存为推理所需的格式。
+
+``` py linenums="235"
+--8<--
+examples/amgnet/amgnet_cylinder.py:235:256
+--8<--
+```
+
+通过运行以下命令，即可执行导出：
+
+```bash
+python amgnet_cylinder.py mode=export
+```
+
+导出的模型将包含 `amgnet_cylinder.pdmodel` (模型结构) 和 `amgnet_cylinder.pdiparams` (模型权重) 文件，保存在配置文件 `INFER.export_path` 所指定的目录中。
+
+#### 3.9.2 创建推理器
+
+为了执行推理，我们创建了一个专用的 `AMGNPredictor` 类，存放于 `deploy/python_infer/amgn_predictor.py`。这个类继承自 `ppsci.deploy.base_predictor.Predictor`，并实现了加载模型和执行预测的核心逻辑。
+
+``` py linenums="28"
+--8<--
+examples/amgnet/deploy/python_infer/amgn_predictor.py:28:87
+--8<--
+```
+
+#### 3.9.3 执行推理
+
+最后，我们实现 `inference` 函数。该函数会实例化 `AMGNPredictor`，加载数据，并循环执行预测，最后将结果可视化。
+
+``` py linenums="259"
+--8<--
+examples/amgnet/amgnet_cylinder.py:259:298
+--8<--
+```
+
+通过以下命令来运行推理：
+
+```bash
+python amgnet_cylinder.py mode=infer
+```
+
+#### 3.9.4 新增配置
+
+为了支持以上功能，需要在 `conf/amgnet_cylinder.yaml` 中添加 `INFER` 配置项。
+
+``` yaml linenums="65"
+--8<--
+examples/amgnet/conf/amgnet_cylinder.yaml:65:68
+--8<--
+```
+
 ## 4. 完整代码
 
 === "airfoil"

examples/amgnet/amgnet_cylinder.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+# 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.
@@ -20,15 +20,17 @@
 from typing import List
 
 import hydra
+import paddle
+import paddle.nn as nn
 import utils
 from omegaconf import DictConfig
 from paddle.nn import functional as F
+from paddle.static import InputSpec
 
 import ppsci
 from ppsci.utils import logger
 
 if TYPE_CHECKING:
-    import paddle
     import pgl
 
 
@@ -212,14 +214,197 @@ def evaluate(cfg: DictConfig):
             )
 
 
+def export(cfg: DictConfig):
+    """Export the model for inference."""
+    # set model
+    model = ppsci.arch.AMGNet(**cfg.MODEL)
+
+    # initialize solver
+    solver = ppsci.solver.Solver(
+        model,
+        pretrained_model_path=cfg.INFER.pretrained_model_path,
+    )
+
+    # Create simplified export model
+    class SimpleExportModel(nn.Layer):
+        def __init__(self, original_model):
+            super(SimpleExportModel, self).__init__()
+            self.output_keys = original_model.output_keys
+
+            # Extract core components from original model
+            self.node_encoder = original_model.encoder.node_model
+            self.post_processor = original_model.post_processor
+            self.decoder = original_model.decoder
+
+        def forward(self, node_feat):
+            """Forward pass that avoids PGL dependencies"""
+            # 1. Node feature encoding
+            encoded_features = self.node_encoder(node_feat)
+
+            # 2. Apply post-processing
+            processed_features = self.post_processor(encoded_features)
+
+            # 3. Apply decoder to get final output
+            output = self.decoder(processed_features)
+
+            return {self.output_keys[0]: output}
+
+    # Create export model
+    export_model = SimpleExportModel(model)
+
+    # Configure export options
+    input_spec = [
+        InputSpec(shape=[None, cfg.MODEL.input_dim], dtype="float32", name="node_feat"),
+    ]
+
+    # Export model
+    solver.export(input_spec, cfg.INFER.export_path, to_func=export_model.forward)
+
+
+def infer(cfg: DictConfig):
+    """Infer using the trained model."""
+    import os
+
+    import amgnet_predictor
+    import numpy as np
+
+    # Create data loader
+    _, dataset = utils.create_dataset(cfg)
+    logger.message("Building dataset...")
+    logger.message("Dataset created successfully")
+
+    # Create AMGNPredictor
+    logger.message("Getting first sample from dataset...")
+    sample = dataset[0]
+
+    # Debug dataset structure
+    logger.message(f"Sample type: {type(sample)}")
+    if isinstance(sample, tuple) and len(sample) >= 3:
+        input_data, label_data, meta = sample
+        logger.message(f"Input data type: {type(input_data)}")
+        logger.message(f"Label data type: {type(label_data)}")
+        logger.message(f"Meta data type: {type(meta)}")
+
+        if isinstance(input_data, dict):
+            for k, v in input_data.items():
+                logger.message(f"Input key: {k}, value type: {type(v)}")
+                if hasattr(v, "x"):
+                    logger.message(f"  Has attribute x: {type(v.x)}")
+                elif isinstance(v, np.ndarray):
+                    logger.message(f"  Is numpy array with shape: {v.shape}")
+    else:
+        logger.message(f"Unexpected sample structure: {sample}")
+        return
+
+    # Extract node features safely
+    try:
+        graph = input_data["input"]
+        logger.message(f"Graph type: {type(graph)}")
+
+        # Handle different types of graph.x
+        if hasattr(graph, "x"):
+            node_feat = graph.x
+            if hasattr(node_feat, "numpy"):
+                node_feat = node_feat.numpy()
+            logger.message(f"Node features shape: {node_feat.shape}")
+        else:
+            # If graph is already a numpy array
+            node_feat = graph
+            logger.message(f"Using graph directly as node features: {node_feat.shape}")
+
+        # Handle different types of label data
+        if isinstance(label_data, dict) and "label" in label_data:
+            label = label_data["label"]
+            if hasattr(label, "y"):
+                ground_truth = label.y
+                if hasattr(ground_truth, "numpy"):
+                    ground_truth = ground_truth.numpy()
+            else:
+                ground_truth = label
+            logger.message(f"Ground truth shape: {ground_truth.shape}")
+        else:
+            logger.message("Could not extract ground truth data")
+            return
+
+        # Handle different types of coordinates
+        if hasattr(graph, "pos"):
+            coords = graph.pos
+            if hasattr(coords, "numpy"):
+                coords = coords.numpy()
+            logger.message(f"Coordinates shape: {coords.shape}")
+        else:
+            # Use first two columns of node_feat as coordinates if available
+            if node_feat.shape[1] >= 2:
+                coords = node_feat[:, :2]
+                logger.message(
+                    f"Using first two columns of node_feat as coordinates: {coords.shape}"
+                )
+            else:
+                # Generate dummy coordinates
+                coords = np.zeros((node_feat.shape[0], 2))
+                logger.message("Using dummy coordinates")
+    except Exception as e:
+        logger.message(f"Error extracting features: {e}")
+        import traceback
+
+        logger.message(traceback.format_exc())
+        return
+
+    # Create directory for result
+    os.makedirs("./result/image/cylinder_infer", exist_ok=True)
+
+    # Initialize predictor
+    predictor = amgnet_predictor.AMGNPredictor(
+        model_path=cfg.INFER.export_path,
+        config_params={
+            "use_mkldnn": cfg.INFER.use_mkldnn,
+            "ir_optim": cfg.INFER.ir_optim,
+        },
+        verbose=True,
+    )
+
+    logger.message("Running inference on sample...")
+
+    # Run inference
+    output = predictor.predict({"node_feat": node_feat})
+
+    logger.message("Generating visualization...")
+
+    # Compare prediction with ground truth
+    pred_result = output["pred"]
+
+    # Extract elements list if available
+    elems_list = None
+    if hasattr(meta, "get"):
+        elems_list = meta.get("elems_list", None)
+
+    # Visualize using the original method
+    utils.log_images(
+        coords,
+        pred_result,
+        ground_truth,
+        elems_list,
+        0,  # Sample index
+        "cylinder_infer",
+    )
+
+    logger.message("Visualization saved to ./result/image/cylinder_infer")
+
+
 @hydra.main(version_base=None, config_path="./conf", config_name="amgnet_cylinder.yaml")
 def main(cfg: DictConfig):
     if cfg.mode == "train":
         train(cfg)
     elif cfg.mode == "eval":
         evaluate(cfg)
+    elif cfg.mode == "export":
+        export(cfg)
+    elif cfg.mode == "infer":
+        infer(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__":