diff --git a/.gitignore b/.gitignore index 491f881f32..c7d2dc9ee9 100644 --- a/.gitignore +++ b/.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_*/ diff --git a/docs/zh/examples/amgnet.md b/docs/zh/examples/amgnet.md index 75d6fa0125..62fec7dbf8 100644 --- a/docs/zh/examples/amgnet.md +++ b/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
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" diff --git a/examples/amgnet/amgnet_cylinder.py b/examples/amgnet/amgnet_cylinder.py index 4203f6052e..a15e9a878d 100644 --- a/examples/amgnet/amgnet_cylinder.py +++ b/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__": diff --git a/examples/amgnet/amgnet_predictor.py b/examples/amgnet/amgnet_predictor.py new file mode 100644 index 0000000000..8106e6387d --- /dev/null +++ b/examples/amgnet/amgnet_predictor.py @@ -0,0 +1,271 @@ +# 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 __future__ import annotations + +import os +import time +from typing import Dict +from typing import List +from typing import Optional + +import numpy as np +from matplotlib import pyplot as plt + +from ppsci.deploy.base_predictor import Predictor +from ppsci.utils import logger + + +class AMGNPredictor(Predictor): + """AMGNet model predictor for inference. + + Supports model performance analysis and result visualization. + """ + + def __init__( + self, + model_path: str, + config_params: Optional[Dict] = None, + verbose: bool = False, + ): + """Initialize AMGNet predictor. + + Args: + model_path: Path to model files. + config_params: Optional configuration parameters. + verbose: Whether to print detailed information. + """ + # Create a minimal config for base class initialization + from omegaconf import OmegaConf + + cfg = OmegaConf.create( + { + "MODEL": {"output_keys": ("pred",)}, + "INFER": { + "export_path": model_path, + "enable_mkldnn": config_params.get("use_mkldnn", True) + if config_params + else True, + "enable_memory_optim": False, # Disable to avoid compatibility issues + "enable_ir_optim": config_params.get("ir_optim", True) + if config_params + else True, + "cpu_threads": 4, + }, + } + ) + + logger.message("Initializing AMGNPredictor...") + super().__init__(cfg) + self.verbose = verbose + + # Initialize input handles for easier access + self.input_handles = [ + self.predictor.get_input_handle(name) + for name in self.predictor.get_input_names() + ] + + logger.message("AMGNPredictor initialized successfully.") + + def predict( + self, inputs: Dict[str, np.ndarray], batch_size: int = 1 + ) -> Dict[str, np.ndarray]: + """Run prediction with AMGNet model. + + Args: + inputs: Input tensor dictionary with keys matching model input names. + batch_size: Batch size for inference. + + Returns: + Output tensor dictionary. + """ + if self.verbose: + logger.message(f"Node feature shape: {inputs['node_feat'].shape}") + + # Warm up inference engine + logger.message("Warming up inference engine...") + self._run_inference(inputs) + + # Measure inference time + start_time = time.time() + output = self._run_inference(inputs) + inference_time = time.time() - start_time + + node_count = inputs["node_feat"].shape[0] + logger.message(f"Inference time: {inference_time:.4f} seconds") + logger.message(f"Nodes processed: {node_count}") + logger.message( + f"Processing speed: {node_count / inference_time:.2f} nodes/second" + ) + + logger.message(f"Prediction successful: {list(output.keys())}") + return output + + def _run_inference(self, inputs: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]: + """Internal method to run actual inference. + + Args: + inputs: Input tensor dictionary. + + Returns: + Output tensor dictionary. + """ + results = {} + + # Set input data + input_names = self.predictor.get_input_names() + for i, name in enumerate(input_names): + if i == 0: # Assume first input is node_feat + handle = self.predictor.get_input_handle(name) + handle.copy_from_cpu(inputs["node_feat"]) + + # Run inference + self.predictor.run() + + # Get output + output_names = self.predictor.get_output_names() + output_handle = self.predictor.get_output_handle(output_names[0]) + output_data = output_handle.copy_to_cpu() + + # Use key 'pred' for consistency with original model + results["pred"] = output_data + + return results + + def input_tensor(self, handle, data): + """Helper method to copy data to input handle. + + Args: + handle: Paddle inference input handle. + data: Input data as numpy array. + + Returns: + The input handle after copying data. + """ + if not isinstance(data, np.ndarray): + data = np.array(data) + handle.copy_from_cpu(data) + return handle + + def analyze_results(self, prediction: np.ndarray, ground_truth: np.ndarray): + """Analyze prediction results compared to ground truth. + + Args: + prediction: Predicted values. + ground_truth: Ground truth values. + """ + # Calculate error metrics + abs_error = np.abs(prediction - ground_truth) + rel_error = abs_error / (np.abs(ground_truth) + 1e-10) + + # Print summary statistics + logger.message("\n===== Result Analysis =====") + logger.message(f"Mean Absolute Error: {np.mean(abs_error):.6f}") + logger.message(f"Max Absolute Error: {np.max(abs_error):.6f}") + logger.message(f"Mean Relative Error: {np.mean(rel_error):.6f}") + logger.message(f"Max Relative Error: {np.max(rel_error):.6f}") + + # Calculate per-channel statistics + for i in range(prediction.shape[1]): + channel_abs_error = np.abs(prediction[:, i] - ground_truth[:, i]) + channel_rel_error = channel_abs_error / (np.abs(ground_truth[:, i]) + 1e-10) + logger.message(f"\nChannel {i} statistics:") + logger.message(f" Mean Absolute Error: {np.mean(channel_abs_error):.6f}") + logger.message(f" Max Absolute Error: {np.max(channel_abs_error):.6f}") + logger.message(f" Mean Relative Error: {np.mean(channel_rel_error):.6f}") + logger.message(f" Max Relative Error: {np.max(channel_rel_error):.6f}") + + # Error histograms + self._plot_error_histogram(abs_error, "Absolute Error") + self._plot_error_histogram(rel_error, "Relative Error") + + def _plot_error_histogram(self, error_data: np.ndarray, title: str, bins: int = 50): + """Plot histogram of errors. + + Args: + error_data: Error data to plot. + title: Plot title. + bins: Number of histogram bins. + """ + plt.figure(figsize=(10, 6)) + plt.hist(error_data.flatten(), bins=bins) + plt.title(f"{title} Distribution") + plt.xlabel(title) + plt.ylabel("Count") + plt.grid(True, alpha=0.3) + plt.tight_layout() + + # Save plot + os.makedirs("./result/analysis", exist_ok=True) + plt.savefig( + f"./result/analysis/{title.lower().replace(' ', '_')}_histogram.png" + ) + plt.close() + + def visualize_comparison( + self, + prediction: np.ndarray, + ground_truth: np.ndarray, + coords: np.ndarray, + channel_names: List[str] = None, + ): + """Visualize comparison between prediction and ground truth. + + Args: + prediction: Predicted values. + ground_truth: Ground truth values. + coords: Coordinates for each node. + channel_names: Names for each channel. + """ + if channel_names is None: + channel_names = [f"Channel_{i}" for i in range(prediction.shape[1])] + + # Create visualization directory + os.makedirs("./result/comparison", exist_ok=True) + + # Plot each channel + for i in range(prediction.shape[1]): + plt.figure(figsize=(15, 7)) + + # Plot ground truth + plt.subplot(1, 3, 1) + sc = plt.scatter( + coords[:, 0], coords[:, 1], c=ground_truth[:, i], cmap="viridis", s=5 + ) + plt.colorbar(sc) + plt.title(f"Ground Truth - {channel_names[i]}") + plt.grid(True, alpha=0.3) + + # Plot prediction + plt.subplot(1, 3, 2) + sc = plt.scatter( + coords[:, 0], coords[:, 1], c=prediction[:, i], cmap="viridis", s=5 + ) + plt.colorbar(sc) + plt.title(f"Prediction - {channel_names[i]}") + plt.grid(True, alpha=0.3) + + # Plot absolute error + plt.subplot(1, 3, 3) + error = np.abs(prediction[:, i] - ground_truth[:, i]) + sc = plt.scatter(coords[:, 0], coords[:, 1], c=error, cmap="hot", s=5) + plt.colorbar(sc) + plt.title(f"Absolute Error - {channel_names[i]}") + plt.grid(True, alpha=0.3) + + plt.tight_layout() + plt.savefig(f"./result/comparison/{channel_names[i]}_comparison.png") + plt.close() + + logger.message("Visualizations saved to ./result/comparison") diff --git a/examples/amgnet/conf/amgnet_cylinder.yaml b/examples/amgnet/conf/amgnet_cylinder.yaml index 39cc559372..e8f7be62a2 100644 --- a/examples/amgnet/conf/amgnet_cylinder.yaml +++ b/examples/amgnet/conf/amgnet_cylinder.yaml @@ -63,5 +63,21 @@ TRAIN: # evaluation settings EVAL: batch_size: 1 - pretrained_model_path: null + # NOTE: The following path is a placeholder, please replace it with your own model path + pretrained_model_path: https://paddle-org.bj.bcebos.com/paddlescience/models/amgnet/amgnet_cylinder_pretrained.pdparams eval_with_no_grad: true + +# inference +INFER: + # Model export settings + export_path: "./inference/amgnet_cylinder" + pretrained_model_path: "/root/.paddlesci/weights/amgnet_cylinder_pretrained.pdparams" + + # Performance settings + use_mkldnn: true + ir_optim: true + + # Analysis settings + analyze_results: true + visualize_comparison: true + channel_names: ["p", "vx", "vy"] diff --git a/examples/amgnet/utils.py b/examples/amgnet/utils.py index 5a8cd3e7f5..e783db58e1 100644 --- a/examples/amgnet/utils.py +++ b/examples/amgnet/utils.py @@ -30,7 +30,6 @@ import matplotlib.pyplot as plt import numpy as np import paddle -from paddle.vision import transforms as T from PIL import Image matplotlib.use("Agg") @@ -140,133 +139,112 @@ def save_image( im.save(fp, format=format) -def log_images( - nodes, - pred, - true, - elems_list, - index, - mode, - aoa=0, - mach=0, - file="field.png", -): - for field in range(pred.shape[1]): - true_img = plot_field( - nodes, - elems_list, - true[:, field], - mode=mode, - col=field, - clim=(-0.8, 0.8), - title="true", - ) - true_img = T.ToTensor()(true_img) - - pred_img = plot_field( - nodes, - elems_list, - pred[:, field], - mode=mode, - col=field, - clim=(-0.8, 0.8), - title="pred", - ) - pred_img = T.ToTensor()(pred_img) - imgs = [pred_img, true_img] - grid = make_grid(paddle.stack(imgs), padding=0) - out_file = file + f"{field}" - if mode == "airfoil": - if aoa == 8.0 and mach == 0.65: - save_image( - grid, "./result/image/" + str(index) + out_file + "_field.png" - ) - save_image( - grid, "./result/image/airfoil/" + str(index) + out_file + "_field.png" - ) - elif mode == "cylinder": - if aoa == 39.0: - save_image( - grid, "./result/image/" + str(index) + out_file + "_field.png" - ) - save_image( - grid, "./result/image/cylinder/" + str(index) + out_file + "_field.png" - ) - else: - raise ValueError( - f"Argument 'mode' should be 'airfoil' or 'cylinder', but got {mode}." - ) - - -def plot_field( - nodes: paddle.Tensor, - elems_list, - field: paddle.Tensor, - mode, - col, - contour=False, - clim=None, - zoom=True, - get_array=True, - out_file=None, - show=False, - title="", -): - elems_list = sum(elems_list, []) - tris, _ = quad2tri(elems_list) - tris = np.array(tris) - x, y = nodes[:, :2].t().detach().numpy() - field = field.detach().numpy() - fig = plt.figure(dpi=800) - if contour: - plt.tricontourf(x, y, tris, field) +def log_images(nodes, pred_field, true_field, elems_list, idx, mode="cylinder"): + """Log images for visualization.""" + import os + + import matplotlib.pyplot as plt + + # 确保结果目录存在 + result_dir = os.path.join("./result/image", mode) + os.makedirs(result_dir, exist_ok=True) + + # 如果是paddle.Tensor，转换为numpy + if not isinstance(pred_field, np.ndarray): + pred_field = pred_field.numpy() + if not isinstance(true_field, np.ndarray): + true_field = true_field.numpy() + + # 压力场 + p_true = true_field[:, 0] + p_pred = pred_field[:, 0] + + # 速度x分量 + vx_true = true_field[:, 1] + vx_pred = pred_field[:, 1] + + # 速度y分量 + vy_true = true_field[:, 2] + vy_pred = pred_field[:, 2] + + # 处理图形数据 + elems = sum(elems_list, []) if elems_list else None + + # 保存压力场图像 + fig_p_true = plot_field(nodes, p_true, elems) + plt.title("True Pressure") + plt.savefig(os.path.join(result_dir, f"p_true_{idx}.png")) + plt.close(fig_p_true) + + fig_p_pred = plot_field(nodes, p_pred, elems) + plt.title("Predicted Pressure") + plt.savefig(os.path.join(result_dir, f"p_pred_{idx}.png")) + plt.close(fig_p_pred) + + # 保存x方向速度场图像 + fig_vx_true = plot_field(nodes, vx_true, elems) + plt.title("True X-Velocity") + plt.savefig(os.path.join(result_dir, f"vx_true_{idx}.png")) + plt.close(fig_vx_true) + + fig_vx_pred = plot_field(nodes, vx_pred, elems) + plt.title("Predicted X-Velocity") + plt.savefig(os.path.join(result_dir, f"vx_pred_{idx}.png")) + plt.close(fig_vx_pred) + + # 保存y方向速度场图像 + fig_vy_true = plot_field(nodes, vy_true, elems) + plt.title("True Y-Velocity") + plt.savefig(os.path.join(result_dir, f"vy_true_{idx}.png")) + plt.close(fig_vy_true) + + fig_vy_pred = plot_field(nodes, vy_pred, elems) + plt.title("Predicted Y-Velocity") + plt.savefig(os.path.join(result_dir, f"vy_pred_{idx}.png")) + plt.close(fig_vy_pred) + + print(f"Saved visualization to {result_dir}") + + +def plot_field(nodes, field, elems, vmin=None, vmax=None): + """Plot heatmap for scalar field in nodes.""" + fig = plt.figure(figsize=(20, 10)) + + # 检查nodes的类型并适当处理 + if isinstance(nodes, np.ndarray): + x, y = nodes[:, 0], nodes[:, 1] # 直接获取x和y坐标 else: - plt.tripcolor(x, y, tris, field) - if clim: - plt.clim(*clim) - colorbar = plt.colorbar() - if mode == "airfoil": - if col == 0: - colorbar.set_label("x-velocity", fontsize=16) - elif col == 1: - colorbar.set_label("pressure", fontsize=16) - elif col == 2: - colorbar.set_label("y-velocity", fontsize=16) - if mode == "cylinder": - if col == 0: - colorbar.set_label("pressure", fontsize=16) - elif col == 1: - colorbar.set_label("x-velocity", fontsize=16) - elif col == 2: - colorbar.set_label("y-velocity", fontsize=16) - if zoom: - if mode == "airfoil": - plt.xlim(left=-0.5, right=1.5) - plt.ylim(bottom=-0.5, top=0.5) - else: - plt.xlim(left=-5, right=5.0) - plt.ylim(bottom=-5, top=5.0) + # 如果是paddle.Tensor或其他支持t()的对象 + x, y = nodes[:, :2].t().detach().numpy() - if title: - plt.title(title) + # 检查field的类型并适当处理 + if isinstance(field, np.ndarray): + value = field + else: + # 如果是paddle.Tensor或其他需要detach的对象 + value = field.detach().numpy() - if out_file is not None: - plt.savefig(out_file) - plt.close() + # 创建三角形单元格 + triangles = [] + if elems is not None: + for cell in elems: + if len(cell) == 3: + triangles.append([cell[0], cell[1], cell[2]]) + elif len(cell) == 4: + triangles.append([cell[0], cell[1], cell[2]]) + triangles.append([cell[0], cell[2], cell[3]]) - if show: - plt.show() - - if get_array: - if mode == "airfoil": - plt.gca().invert_yaxis() - fig.canvas.draw() - array = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8) - array = array.reshape(fig.canvas.get_width_height()[::-1] + (3,)) - fig.clf() - fig.clear() - plt.close() - return array + if len(triangles) > 0: + plt.tricontourf(x, y, triangles, value, 100, cmap="jet", vmin=vmin, vmax=vmax) + else: + plt.tricontourf(x, y, value, 100, cmap="jet", vmin=vmin, vmax=vmax) + plt.colorbar() + plt.axes().set_aspect("equal") + plt.xlabel("x") + plt.ylabel("y") + plt.tight_layout() + + return fig def quad2tri(elems): @@ -285,3 +263,97 @@ def quad2tri(elems): else paddle.to_tensor([], dtype=paddle.int64) ) return new_elems, new_edges + + +def create_dataset(cfg): + """Create dataset for inference. + + Args: + cfg: Configuration object. + + Returns: + data_loader, dataset: Tuple containing data loader and dataset. + """ + from ppsci.data import dataset + + # Create dataset + eval_dataset = dataset.MeshCylinderDataset( + input_keys=("input",), + label_keys=("label",), + data_dir=cfg.EVAL_DATA_DIR, + mesh_graph_path=cfg.EVAL_MESH_GRAPH_PATH, + ) + + return None, eval_dataset + + +def visualize_result(data, save_path, channel_names=None): + """Visualize prediction results. + + Args: + data: Dictionary containing 'pred' and 'true' fields. + save_path: Path to save visualizations. + channel_names: List of channel names. + """ + if channel_names is None: + channel_names = ["p", "vx", "vy"] + + # Ensure directory exists + os.makedirs(save_path, exist_ok=True) + + # Get prediction and ground truth + pred = data["pred"] + true = data["true"] + + # Visualize each channel + for i, name in enumerate(channel_names): + if i >= pred.shape[1]: + continue + + # Plot ground truth heatmap + plt.figure(figsize=(10, 8)) + plt.imshow(true[:, i].reshape(-1, 1), cmap="viridis", aspect="auto") + plt.colorbar() + plt.title(f"True {name}") + plt.savefig(os.path.join(save_path, f"{name}_true_0.png")) + plt.close() + + # Plot prediction heatmap + plt.figure(figsize=(10, 8)) + plt.imshow(pred[:, i].reshape(-1, 1), cmap="viridis", aspect="auto") + plt.colorbar() + plt.title(f"Predicted {name}") + plt.savefig(os.path.join(save_path, f"{name}_pred_0.png")) + plt.close() + + +def graph_collate_fn(batch): + """自定义的collate函数，用于处理包含图对象的批次数据。 + + Args: + batch: 批次数据，每个元素是一个样本 + + Returns: + 处理后的批次数据 + """ + input_dict = {} + label_dict = {} + meta_dict = {} + + # 对于每个样本 + for sample in batch: + for k, v in sample[0].items(): + if k not in input_dict: + input_dict[k] = [] + input_dict[k].append(v) + + for k, v in sample[1].items(): + if k not in label_dict: + label_dict[k] = [] + label_dict[k].append(v) + + meta_dict_sample = sample[2] + if not meta_dict: + meta_dict = meta_dict_sample + + return input_dict, label_dict, meta_dict diff --git a/ppsci/__init__.py b/ppsci/__init__.py index bde7aa49a5..918beba1e3 100644 --- a/ppsci/__init__.py +++ b/ppsci/__init__.py @@ -16,6 +16,7 @@ from ppsci import autodiff # isort:skip from ppsci import constraint # isort:skip from ppsci import data # isort:skip +from ppsci import deploy # isort:skip from ppsci import equation # isort:skip from ppsci import geometry # isort:skip from ppsci import loss # isort:skip @@ -45,6 +46,7 @@ "autodiff", "constraint", "data", + "deploy", "equation", "geometry", "loss", diff --git a/ppsci/arch/amgnet.py b/ppsci/arch/amgnet.py index ce728317d6..9a7c868af0 100644 --- a/ppsci/arch/amgnet.py +++ b/ppsci/arch/amgnet.py @@ -628,6 +628,34 @@ def forward(self, x: Dict[str, "pgl.Graph"]) -> Dict[str, paddle.Tensor]: pred_field = self.decoder(node_features) return {self.output_keys[0]: pred_field} + def forward_export(self, node_feat, edge_feat): + """Forward pass for export mode. + + Simplified forward function designed for exported model to avoid direct PGL usage. + + Args: + node_feat: Node feature tensor. + edge_feat: Edge feature tensor. + + Returns: + Dict[str, paddle.Tensor]: Dictionary containing prediction results. + """ + # Create PGL graph with provided features + import pgl + + graph = pgl.Graph(num_nodes=node_feat.shape[0], edges=[[0, 0]]) # Dummy edge + graph.x = node_feat + graph.edge_attr = edge_feat + graph.pos = node_feat[:, :2] # Use first two dims as position + graph.edge_index = paddle.to_tensor([[0], [0]]) # Dummy edge index + + # Process with regular forward + latent_graph = self.encoder(graph) + x, p = self.processor(latent_graph, speed=self.speed) + node_features = self._spa_compute(x, p) + pred_field = self.decoder(node_features) + return {self.output_keys[0]: pred_field} + def _make_mlp(self, output_dim: int, input_dim: int = 5, layer_norm: bool = True): widths = (self._latent_dim,) * self._num_layers + (output_dim,) network = FullyConnectedLayer(input_dim, widths) diff --git a/ppsci/deploy/__init__.py b/ppsci/deploy/__init__.py new file mode 100644 index 0000000000..9ba2638d5a --- /dev/null +++ b/ppsci/deploy/__init__.py @@ -0,0 +1,19 @@ +# Copyright (c) 2023 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 __future__ import annotations + +from ppsci.deploy.base_predictor import Predictor + +__all__ = ["Predictor"] diff --git a/ppsci/deploy/base_predictor.py b/ppsci/deploy/base_predictor.py new file mode 100644 index 0000000000..39a3b189f3 --- /dev/null +++ b/ppsci/deploy/base_predictor.py @@ -0,0 +1,122 @@ +# Copyright (c) 2023 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 __future__ import annotations + +import os +from typing import Dict + +import paddle +import paddle.inference as paddle_infer +from omegaconf import DictConfig + +from ppsci.utils import logger + + +class Predictor: + """Base class for model predictors. + + Args: + cfg (DictConfig): Configuration object. + """ + + def __init__(self, cfg: DictConfig): + self.cfg = cfg + self.output_keys = cfg.MODEL.output_keys + self.enable_mkldnn = getattr(cfg.INFER, "enable_mkldnn", False) + self.enable_memory_optim = getattr(cfg.INFER, "enable_memory_optim", True) + self.enable_ir_optim = getattr(cfg.INFER, "enable_ir_optim", True) + self.cpu_threads = getattr(cfg.INFER, "cpu_threads", 4) + + # Load model + self._load_model() + + def _load_model(self): + """Load model from exported files.""" + model_dir = os.path.dirname(self.cfg.INFER.export_path) + model_prefix = os.path.basename(self.cfg.INFER.export_path) + + logger.message(f"Loading model from {self.cfg.INFER.export_path}") + + # Create config + config = paddle_infer.Config() + + # Check if model files exist + if paddle.framework.use_pir_api(): + model_file = os.path.join(model_dir, model_prefix + ".json") + params_file = os.path.join(model_dir, model_prefix + ".pdiparams") + if not os.path.exists(model_file) or not os.path.exists(params_file): + raise ValueError(f"Model files not found: {model_file}, {params_file}") + # For PIR API + config.set_prog_file(model_file) # Just use set_prog_file in both cases + config.set_params_file(params_file) + else: + model_file = os.path.join(model_dir, model_prefix + ".pdmodel") + params_file = os.path.join(model_dir, model_prefix + ".pdiparams") + if not os.path.exists(model_file) or not os.path.exists(params_file): + raise ValueError(f"Model files not found: {model_file}, {params_file}") + config.set_prog_file(model_file) + config.set_params_file(params_file) + + # CPU configurations - expand with more optimization options + config.disable_gpu() + config.set_cpu_math_library_num_threads(self.cpu_threads) + + # 添加性能优化配置 (安全的API调用方式) + if self.enable_mkldnn: + try: + if hasattr(config, "enable_mkldnn"): + config.enable_mkldnn() + logger.message("MKL-DNN acceleration enabled") + except Exception as e: + logger.warning(f"Failed to enable MKL-DNN: {e}") + + if self.enable_memory_optim: + try: + if hasattr(config, "enable_memory_optim"): + config.enable_memory_optim() + logger.message("Memory optimization enabled") + except Exception as e: + logger.warning(f"Failed to enable memory optimization: {e}") + + if self.enable_ir_optim: + try: + if hasattr(config, "switch_ir_optimize") or hasattr( + config, "switch_ir_optim" + ): + # 不同版本API名称可能不同 + if hasattr(config, "switch_ir_optimize"): + config.switch_ir_optimize(True) + elif hasattr(config, "switch_ir_optim"): + config.switch_ir_optim(True) + logger.message("IR optimization enabled") + except Exception as e: + logger.warning(f"Failed to enable IR optimization: {e}") + + # Create predictor with enhanced configuration + self.predictor = paddle_infer.create_predictor(config) + logger.message("Predictor created successfully with optimized configuration") + + def predict(self, input_dict: Dict, batch_size: int = 64) -> Dict: + """Predicts the output of the model for a given input. + + Args: + input_dict (Dict): Input data. + batch_size (int, optional): Batch size for prediction. Defaults to 64. + + Returns: + Dict: Predicted output. + """ + # This is a base method, should be implemented by derived classes + raise NotImplementedError diff --git a/ppsci/solver/solver.py b/ppsci/solver/solver.py index ffe7ef291b..a2319685d7 100644 --- a/ppsci/solver/solver.py +++ b/ppsci/solver/solver.py @@ -893,19 +893,20 @@ def predict( @misc.run_on_eval_mode def export( self, - input_spec: List[Dict[str, InputSpec]], + input_spec: List[Union[Dict[str, InputSpec], InputSpec]], export_path: str, with_onnx: bool = False, skip_prune_program: bool = False, *, full_graph: bool = True, ignore_modules: Optional[List[ModuleType]] = None, + to_func: Optional[Callable] = None, ): """ Convert model to static graph model and export to files. Args: - input_spec (List[Dict[str, InputSpec]]): InputSpec describes the signature + input_spec (List[Union[Dict[str, InputSpec], InputSpec]]): InputSpec describes the signature information of the model input. export_path (str): The path prefix to save model. with_onnx (bool, optional): Whether to export model into onnx after @@ -919,6 +920,8 @@ def export( conversion. Builtin modules that have been ignored are collections, pdb, copy, inspect, re, numpy, logging, six. For example, einops can be added here. Defaults to None. + to_func (Optional[Callable]): Function to be exported, if provided will be used instead of model.forward. + Defaults to None. """ if ignore_modules is not None: jit.ignore_module(ignore_modules) @@ -931,9 +934,18 @@ def export( "model will be random initialized." ) + # Check if to_func is provided + if to_func is not None: + forward_func = to_func + # Check if model has a forward_export method + elif hasattr(self.model, "forward_export"): + forward_func = self.model.forward_export + else: + forward_func = self.model.forward + # convert model to static graph model static_model = jit.to_static( - self.model, + forward_func, input_spec=input_spec, full_graph=full_graph, )