Merge branch 'develop' into clean_code

HydrogenSulfate · web-flow · commit f4cdcd4ebbc0 · 2024-05-17T14:16:07.000+08:00
diff --git a/docs/zh/examples/darcy2d.md b/docs/zh/examples/darcy2d.md
@@ -14,6 +14,18 @@
     python darcy2d.py mode=eval EVAL.pretrained_model_path=https://paddle-org.bj.bcebos.com/paddlescience/models/darcy2d/darcy2d_pretrained.pdparams
     ```
 
+=== "模型导出命令"
+
+    ``` sh
+    python darcy2d.py mode=export
+    ```
+
+=== "模型推理命令"
+
+    ``` sh
+    python darcy2d.py mode=infer
+    ```
+
 | 预训练模型  | 指标 |
 |:--| :--|
 | [darcy2d_pretrained.pdparams](https://paddle-org.bj.bcebos.com/paddlescience/models/darcy2d/darcy2d_pretrained.pdparams) | loss(Residual): 0.36500<br>MSE.poisson(Residual): 0.00006 |
diff --git a/docs/zh/examples/deeponet.md b/docs/zh/examples/deeponet.md
@@ -26,6 +26,18 @@
     python deeponet.py mode=eval EVAL.pretrained_model_path=https://paddle-org.bj.bcebos.com/paddlescience/models/deeponet/deeponet_pretrained.pdparams
     ```
 
+=== "模型导出命令"
+
+    ``` sh
+    python deeponet.py mode=export
+    ```
+
+=== "模型推理命令"
+
+    ``` sh
+    python deeponet.py mode=infer
+    ```
+
 | 预训练模型  | 指标 |
 |:--| :--|
 | [deeponet_pretrained.pdparams](https://paddle-org.bj.bcebos.com/paddlescience/models/deeponet/deeponet_pretrained.pdparams) | loss(G_eval): 0.00003<br>L2Rel.G(G_eval): 0.01799 |
diff --git a/docs/zh/examples/hpinns.md b/docs/zh/examples/hpinns.md
@@ -26,6 +26,25 @@
     python holography.py mode=eval EVAL.pretrained_model_path=https://paddle-org.bj.bcebos.com/paddlescience/models/hPINNs/hpinns_pretrained.pdparams
     ```
 
+=== "模型导出命令"
+
+    ``` sh
+    python holography.py mode=export
+    ```
+
+=== "模型推理命令"
+
+    ``` sh
+    # linux
+    wget -nc https://paddle-org.bj.bcebos.com/paddlescience/datasets/hPINNs/hpinns_holo_train.mat -P ./datasets/
+    wget -nc https://paddle-org.bj.bcebos.com/paddlescience/datasets/hPINNs/hpinns_holo_valid.mat -P ./datasets/
+    # windows
+    # curl https://paddle-org.bj.bcebos.com/paddlescience/datasets/hPINNs/hpinns_holo_train.mat --output ./datasets/hpinns_holo_train.mat
+    # curl https://paddle-org.bj.bcebos.com/paddlescience/datasets/hPINNs/hpinns_holo_valid.mat --output ./datasets/hpinns_holo_valid.mat
+    python holography.py mode=infer
+    ```
+
+
 | 预训练模型  | 指标 |
 |:--| :--|
 | [hpinns_pretrained.pdparams](https://paddle-org.bj.bcebos.com/paddlescience/models/hPINNs/hpinns_pretrained.pdparams) | loss(opt_sup): 0.05352<br>MSE.eval_metric(opt_sup): 0.00002<br>loss(val_sup): 0.02205<br>MSE.eval_metric(val_sup): 0.00001 |
diff --git a/examples/darcy/conf/darcy2d.yaml b/examples/darcy/conf/darcy2d.yaml
@@ -35,6 +35,7 @@ hydra:
 mode: train # running mode: train/eval
 seed: 42
 output_dir: ${hydra:run.dir}
+log_freq: 20
 
 # set working condition
 NPOINT_PDE: 9801  # 99 ** 2
@@ -74,3 +75,20 @@ EVAL:
   batch_size:
     residual_validator: 8192
   pretrained_model_path: null
+
+INFER:
+  pretrained_model_path: https://paddle-org.bj.bcebos.com/paddlescience/models/darcy2d/darcy2d_pretrained.pdparams
+  export_path: ./inference/darcy2d
+  pdmodel_path: ${INFER.export_path}.pdmodel
+  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: 8192
+  num_cpu_threads: 10
+  batch_size: 8192
diff --git a/examples/darcy/darcy2d.py b/examples/darcy/darcy2d.py
@@ -275,14 +275,66 @@ def poisson_ref_compute_func(_in):
     solver.visualize()
 
 
+def export(cfg: DictConfig):
+    # set model
+    model = ppsci.arch.MLP(**cfg.MODEL)
+
+    # 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, 1], "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)
+
+    # set geometry
+    geom = {"rect": ppsci.geometry.Rectangle((0.0, 0.0), (1.0, 1.0))}
+    # manually collate input data for visualization,
+    input_dict = geom["rect"].sample_interior(
+        cfg.NPOINT_PDE + cfg.NPOINT_BC, evenly=True
+    )
+    output_dict = predictor.predict(
+        {key: input_dict[key] for key in cfg.MODEL.input_keys}, cfg.INFER.batch_size
+    )
+    # mapping data to cfg.INFER.output_keys
+    output_dict = {
+        store_key: output_dict[infer_key]
+        for store_key, infer_key in zip(cfg.MODEL.output_keys, output_dict.keys())
+    }
+    ppsci.visualize.save_vtu_from_dict(
+        "./visual/darcy2d.vtu",
+        {**input_dict, **output_dict},
+        input_dict.keys(),
+        cfg.MODEL.output_keys,
+    )
+
+
 @hydra.main(version_base=None, config_path="./conf", config_name="darcy2d.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__":
diff --git a/examples/hpinns/conf/hpinns.yaml b/examples/hpinns/conf/hpinns.yaml
@@ -37,6 +37,7 @@ seed: 42
 output_dir: ${hydra:run.dir}
 DATASET_PATH: ./datasets/hpinns_holo_train.mat
 DATASET_PATH_VALID: ./datasets/hpinns_holo_valid.mat
+log_freq: 20
 
 # set working condition
 TRAIN_MODE: aug_lag  # "soft", "penalty", "aug_lag"
@@ -77,3 +78,22 @@ TRAIN:
 # evaluation settings
 EVAL:
   pretrained_model_path: null
+
+# inference settings
+INFER:
+  pretrained_model_path: "https://paddle-org.bj.bcebos.com/paddlescience/models/hPINNs/hpinns_pretrained.pdparams"
+  export_path: ./inference/hpinns
+  pdmodel_path: ${INFER.export_path}.pdmodel
+  pdiparams_path: ${INFER.export_path}.pdiparams
+  output_keys: ["e_re", "e_im", "eps"]
+  device: gpu
+  engine: native
+  precision: fp32
+  onnx_path: ${INFER.export_path}.onnx
+  ir_optim: true
+  min_subgraph_size: 10
+  gpu_mem: 8000
+  gpu_id: 0
+  batch_size: 128
+  max_batch_size: 128
+  num_cpu_threads: 4
diff --git a/examples/hpinns/holography.py b/examples/hpinns/holography.py
@@ -383,14 +383,86 @@ def evaluate(cfg: DictConfig):
     solver.eval()
 
 
+def export(cfg: DictConfig):
+    # set model
+    model_re = ppsci.arch.MLP(**cfg.MODEL.re_net)
+    model_im = ppsci.arch.MLP(**cfg.MODEL.im_net)
+    model_eps = ppsci.arch.MLP(**cfg.MODEL.eps_net)
+
+    # register transform
+    model_re.register_input_transform(func_module.transform_in)
+    model_im.register_input_transform(func_module.transform_in)
+    model_eps.register_input_transform(func_module.transform_in)
+
+    model_re.register_output_transform(func_module.transform_out_real_part)
+    model_im.register_output_transform(func_module.transform_out_imaginary_part)
+    model_eps.register_output_transform(func_module.transform_out_epsilon)
+
+    # wrap to a model_list
+    model_list = ppsci.arch.ModelList((model_re, model_im, model_eps))
+
+    # 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 ["x", "y"]},
+    ]
+    solver.export(input_spec, cfg.INFER.export_path)
+
+
+def inference(cfg: DictConfig):
+    from deploy.python_infer import pinn_predictor
+
+    predictor = pinn_predictor.PINNPredictor(cfg)
+
+    valid_dict = ppsci.utils.reader.load_mat_file(
+        cfg.DATASET_PATH_VALID, ("x_val", "y_val", "bound")
+    )
+    input_dict = {"x": valid_dict["x_val"], "y": valid_dict["y_val"]}
+
+    output_dict = predictor.predict(input_dict, cfg.INFER.batch_size)
+
+    # mapping data to cfg.INFER.output_keys
+    output_dict = {
+        store_key: output_dict[infer_key]
+        for store_key, infer_key in zip(cfg.INFER.output_keys, output_dict.keys())
+    }
+
+    # plotting E and eps
+    N = ((func_module.l_BOX[1] - func_module.l_BOX[0]) / 0.05).astype(int)
+    input_eval = np.stack((input_dict["x"], input_dict["y"]), axis=-1).reshape(
+        N[0], N[1], 2
+    )
+    e_re = output_dict["e_re"].reshape(N[0], N[1])
+    e_im = output_dict["e_im"].reshape(N[0], N[1])
+    eps = output_dict["eps"].reshape(N[0], N[1])
+    v_visual = e_re**2 + e_im**2
+    field_visual = np.stack((v_visual, eps), axis=-1)
+    plot_module.field_name = ["Fig7_E", "Fig7_eps"]
+    plot_module.FIGNAME = "hpinns_pred"
+    plot_module.plot_field_holo(input_eval, field_visual)
+
+
 @hydra.main(version_base=None, config_path="./conf", config_name="hpinns.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__":
diff --git a/examples/hpinns/plotting.py b/examples/hpinns/plotting.py
@@ -20,6 +20,7 @@
 from typing import Callable
 from typing import Dict
 from typing import List
+from typing import Optional
 
 import functions as func_module
 import matplotlib.pyplot as plt
@@ -114,16 +115,16 @@ def prepare_data(solver: ppsci.solver.Solver, expr_dict: Dict[str, Callable]):
 def plot_field_holo(
     coord_visual: np.ndarray,
     field_visual: np.ndarray,
-    coord_lambda: np.ndarray,
-    field_lambda: np.ndarray,
+    coord_lambda: Optional[np.ndarray] = None,
+    field_lambda: Optional[np.ndarray] = None,
 ):
     """Plot fields of of holography example.
 
     Args:
         coord_visual (np.ndarray): The coord of epsilon and |E|**2.
         field_visual (np.ndarray): The filed of epsilon and |E|**2.
-        coord_lambda (np.ndarray): The coord of lambda.
-        field_lambda (np.ndarray): The filed of lambda.
+        coord_lambda (Optional[np.ndarray], optional): The coord of lambda. Defaults to None.
+        field_lambda (Optional[np.ndarray], optional): The filed of lambda. Defaults to None.
     """
     fmin, fmax = np.array([0, 1.0]), np.array([0.6, 12])
     cmin, cmax = coord_visual.min(axis=(0, 1)), coord_visual.max(axis=(0, 1))
@@ -168,7 +169,7 @@ def plot_field_holo(
             cb = plt.colorbar()
             plt.axis((emin[0], emax[0], emin[1], emax[1]))
             plt.clim(vmin=fmin[fi], vmax=fmax[fi])
-        else:
+        elif coord_lambda is not None and field_lambda is not None:
             # Fig_6C_lambda_
             plt.figure(fi * 100 + 101, figsize=(8, 6))
             plt.clf()
diff --git a/examples/operator_learning/conf/deeponet.yaml b/examples/operator_learning/conf/deeponet.yaml
@@ -72,3 +72,21 @@ TRAIN:
 EVAL:
   pretrained_model_path: null
   eval_with_no_grad: true
+
+# inference settings
+INFER:
+  pretrained_model_path: "https://paddle-org.bj.bcebos.com/paddlescience/models/deeponet/deeponet_pretrained.pdparams"
+  export_path: ./inference/deeponet
+  pdmodel_path: ${INFER.export_path}.pdmodel
+  pdiparams_path: ${INFER.export_path}.pdiparams
+  device: gpu
+  engine: native
+  precision: fp32
+  onnx_path: ${INFER.export_path}.onnx
+  ir_optim: true
+  min_subgraph_size: 10
+  gpu_mem: 4000
+  gpu_id: 0
+  max_batch_size: 128
+  num_cpu_threads: 4
+  batch_size: 128
diff --git a/examples/operator_learning/deeponet.py b/examples/operator_learning/deeponet.py
diff --git a/ppsci/geometry/geometry.py b/ppsci/geometry/geometry.py
