【PPSCI Export&Infer No.1】 Add export and inference for deephpms_schrodinger -part

docs/zh/examples/deephpms.md

@@ -133,6 +133,23 @@
     # curl https://paddle-org.bj.bcebos.com/paddlescience/datasets/DeepHPMs/NLS.mat --create-dirs -o ./datasets/NLS.mat
     python schrodinger.py mode=eval DATASET_PATH=./datasets/NLS.mat DATASET_PATH_SOL=./datasets/NLS.mat EVAL.pretrained_model_path=https://paddle-org.bj.bcebos.com/paddlescience/models/DeepHPMs/schrodinger_pretrained.pdparams
     ```
+=== "模型导出命令"
+
+    ``` sh
+    # 案例8
+    # linux
+    wget -nc https://paddle-org.bj.bcebos.com/paddlescience/datasets/DeepHPMs/NLS.mat -P ./datasets/
+    # windows
+    # curl https://paddle-org.bj.bcebos.com/paddlescience/datasets/DeepHPMs/NLS.mat --create-dirs -o ./datasets/NLS.mat
+    python schrodinger.py mode=export INFER.pretrained_model_path=https://paddle-org.bj.bcebos.com/paddlescience/models/DeepHPMs/schrodinger_pretrained.pdparams
+    ```
+
+=== "模型推理命令"
+
+    ``` sh
+    # 案例8
+    python schrodinger.py mode=infer INFER.pretrained_model_path=https://paddle-org.bj.bcebos.com/paddlescience/models/DeepHPMs/schrodinger_pretrained.pdparams
+    ```
 
 | 序号 | 案例名称 | stage1、2 数据集 | stage3(eval)数据集 | 预训练模型 | 指标 |
 | :-- | :-- | :-- | :-- | :-- | :-- |

examples/deephpms/conf/schrodinger.yaml

@@ -75,3 +75,23 @@ TRAIN:
 # evaluation settings
 EVAL:
   pretrained_model_path: null
+
+# inference settings
+INFER:
+  pretrained_model_path: https://paddle-org.bj.bcebos.com/paddlescience/models/DeepHPMs/schrodinger_pretrained.pdparams
+  export_path: ./inference/schrodinger
+  pdmodel_path: ${INFER.export_path}.json
+  pdiparams_path: ${INFER.export_path}.pdiparams
+  onnx_path: ${INFER.export_path}.onnx
+  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
+  batch_size: 1024
+  t_points: 100
+  x_points: 100

examples/deephpms/schrodinger.py

@@ -535,14 +535,115 @@ def transform_fg(_in):
     )
 
 
+def export(cfg: DictConfig):
+    # set model
+    model_idn_u = ppsci.arch.MLP(**cfg.MODEL.idn_u_net)
+    model_idn_v = ppsci.arch.MLP(**cfg.MODEL.idn_v_net)
+
+    # initialize transform
+    def transform_uv(_in):
+        t_lb = paddle.to_tensor(cfg.T_LB)
+        t_ub = paddle.to_tensor(np.pi / cfg.T_UB)
+        x_lb = paddle.to_tensor(cfg.X_LB)
+        x_ub = paddle.to_tensor(cfg.X_UB)
+
+        t, x = _in["t"], _in["x"]
+        t = 2.0 * (t - t_lb) * paddle.pow((t_ub - t_lb), -1) - 1.0
+        x = 2.0 * (x - x_lb) * paddle.pow((x_ub - x_lb), -1) - 1.0
+        input_trans = {"t": t, "x": x}
+        return input_trans
+
+    # register transform
+    model_idn_u.register_input_transform(transform_uv)
+    model_idn_v.register_input_transform(transform_uv)
+
+    # initialize model list
+    model_list = ppsci.arch.ModelList((model_idn_u, model_idn_v))
+
+    # initialize solver
+    solver = ppsci.solver.Solver(
+        model_list,
+        pretrained_model_path=cfg.INFER.pretrained_model_path,
+    )
+
+    # export model
+    from paddle.static import InputSpec
+
+    input_spec = [
+        {key: InputSpec([None, 1], "float32", name=key) for key in ["t", "x"]},
+    ]
+    solver.export(input_spec, cfg.INFER.export_path, with_onnx=False)
+
+
+def inference(cfg: DictConfig):
+    from deploy.python_infer import pinn_predictor
+
+    predictor = pinn_predictor.PINNPredictor(cfg)
+    dataset_val = reader.load_mat_file(
+        cfg.DATASET_PATH_SOL,
+        ("t", "x", "uv_sol", "u_sol", "v_sol"),
+        {
+            "t": "t_ori",
+            "x": "x_ori",
+            "uv_sol": "Exact_uv_ori",
+            "u_sol": "u_star",
+            "v_sol": "v_star",
+        },
+    )
+    t_mesh, x_mesh = np.meshgrid(
+        np.squeeze(dataset_val["t"]), np.squeeze(dataset_val["x"])
+    )
+    input_dict = {
+        "t": t_mesh.flatten()[:, None].astype(np.float32),
+        "x": x_mesh.flatten()[:, None].astype(np.float32),
+    }
+    output_dict = predictor.predict(input_dict, cfg.INFER.batch_size)
+
+    # mapping data to output keys
+    output_dict = {
+        store_key: output_dict[infer_key]
+        for store_key, infer_key in zip(["u_idn", "v_idn"], output_dict.keys())
+    }
+
+    # 计算误差并保存
+    u_pred = output_dict["u_idn"]
+    v_pred = output_dict["v_idn"]
+    uv_pred = np.sqrt(u_pred**2 + v_pred**2)
+
+    error_uv = np.linalg.norm(
+        dataset_val["uv_sol"].flatten() - uv_pred.flatten(), 2
+    ) / np.linalg.norm(dataset_val["uv_sol"].flatten(), 2)
+    logger.info(f"L2 relative error: {error_uv:.6f}")
+
+    np.savez(
+        osp.join(cfg.output_dir, "schrodinger_pred.npz"),
+        **input_dict,
+        **output_dict,
+        uv_pred=uv_pred,
+    )
+
+    ppsci.visualize.save_vtu_from_dict(
+        "./schrodinger_pred.vtu",
+        {**input_dict, **output_dict, "uv_pred": uv_pred},
+        input_dict.keys(),
+        ["u_idn", "v_idn", "uv_pred"],
+    )
+
+
 @hydra.main(version_base=None, config_path="./conf", config_name="schrodinger.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":
+        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__":