add export&infer 26-27 (#893)

smallpoxscattered · web-flow · commit 91dadab99057 · 2024-05-13T15:17:20.000+08:00
diff --git a/docs/zh/examples/control_arm.md b/docs/zh/examples/control_arm.md
@@ -46,6 +46,42 @@
         python inverse_parameter.py mode=eval EVAL.pretrained_model_path=https://paddle-org.bj.bcebos.com/paddlescience/models/control_arm/inverse_x_axis_pretrained.pdparams
         ```
 
+=== "模型导出命令"
+
+    === "正问题：受力分析求解"
+
+        ``` sh
+        python forward_analysis.py mode=export
+        ```
+
+    === "逆问题：参数逆推求解"
+
+        ``` sh
+        python inverse_parameter.py mode=export
+        ```
+
+=== "模型推理命令"
+
+    === "正问题：受力分析求解"
+
+        ``` sh
+        # linux
+        wget -nc https://paddle-org.bj.bcebos.com/paddlescience/datasets/control_arm/control_arm.stl -P ./datasets/
+        # windows
+        # curl https://paddle-org.bj.bcebos.com/paddlescience/datasets/control_arm/control_arm.stl --output ./datasets/control_arm.stl
+        python forward_analysis.py mode=infer
+        ```
+
+    === "逆问题：参数逆推求解"
+
+        ``` sh
+        # linux
+        wget -nc https://paddle-org.bj.bcebos.com/paddlescience/datasets/control_arm/control_arm.stl -P ./datasets/
+        # windows
+        # curl https://paddle-org.bj.bcebos.com/paddlescience/datasets/control_arm/control_arm.stl --output ./datasets/control_arm.stl
+        python inverse_parameter.py mode=infer
+        ```
+
 | 预训练模型  | 指标 |
 |:--| :--|
 | [inverse_x_axis_pretrained.pdparams](https://paddle-org.bj.bcebos.com/paddlescience/models/control_arm/inverse_x_axis_pretrained.pdparams) | loss(geo_eval): 0.02505<br>L2Rel.lambda_(geo_eval): 0.06025<br>L2Rel.mu(geo_eval): 0.07949 |
diff --git a/examples/control_arm/conf/forward_analysis.yaml b/examples/control_arm/conf/forward_analysis.yaml
@@ -11,6 +11,8 @@ hydra:
           - TRAIN.checkpoint_path
           - TRAIN.pretrained_model_path
           - EVAL.pretrained_model_path
+          - INFER.pretrained_model_path
+          - INFER.export_path
           - mode
           - output_dir
           - log_freq
@@ -87,3 +89,21 @@ TRAIN:
 EVAL:
   eval_with_no_grad: true
   pretrained_model_path: null
+
+
+INFER:
+  pretrained_model_path: https://paddle-org.bj.bcebos.com/paddlescience/models/control_arm/forward_x_axis_pretrained.pdparams
+  export_path: ./inference/forward_analysis
+  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: 64
diff --git a/examples/control_arm/conf/inverse_parameter.yaml b/examples/control_arm/conf/inverse_parameter.yaml
@@ -11,6 +11,8 @@ hydra:
           - TRAIN.checkpoint_path
           - TRAIN.pretrained_model_path
           - EVAL.pretrained_model_path
+          - INFER.pretrained_model_path
+          - INFER.export_path
           - mode
           - output_dir
           - log_freq
@@ -99,3 +101,20 @@ EVAL:
   batch_size:
     validator: 1024
     visualizer_vtu: 100000
+
+INFER:
+  pretrained_model_path: https://paddle-org.bj.bcebos.com/paddlescience/models/control_arm/inverse_x_axis_pretrained.pdparams
+  export_path: ./inference/inverse_parameter
+  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: 64
diff --git a/examples/control_arm/forward_analysis.py b/examples/control_arm/forward_analysis.py
@@ -281,6 +281,81 @@ def evaluate(cfg: DictConfig):
     solver.visualize()
 
 
+def export(cfg: DictConfig):
+    from paddle.static import InputSpec
+
+    # set model
+    disp_net = ppsci.arch.MLP(**cfg.MODEL.disp_net)
+    stress_net = ppsci.arch.MLP(**cfg.MODEL.stress_net)
+    # wrap to a model_list
+    model_list = ppsci.arch.ModelList((disp_net, stress_net))
+
+    # load pretrained model
+    solver = ppsci.solver.Solver(
+        model=model_list, pretrained_model_path=cfg.INFER.pretrained_model_path
+    )
+
+    # export models
+    input_spec = [
+        {
+            key: InputSpec([None, 1], "float32", name=key)
+            for key in cfg.MODEL.disp_net.input_keys
+        },
+    ]
+    solver.export(input_spec, cfg.INFER.export_path)
+
+
+def inference(cfg: DictConfig):
+    from deploy.python_infer import pinn_predictor
+    from ppsci.visualize import vtu
+
+    # set model predictor
+    predictor = pinn_predictor.PINNPredictor(cfg)
+
+    # set geometry
+    control_arm = ppsci.geometry.Mesh(cfg.GEOM_PATH)
+    # geometry bool operation
+    geo = control_arm
+    geom = {"geo": geo}
+    # set bounds
+    BOUNDS_X, BOUNDS_Y, BOUNDS_Z = control_arm.bounds
+
+    # set visualizer(optional)
+    # add inferencer data
+    samples = geom["geo"].sample_interior(
+        cfg.TRAIN.batch_size.visualizer_vtu,
+        criteria=lambda x, y, z: (
+            (BOUNDS_X[0] < x)
+            & (x < BOUNDS_X[1])
+            & (BOUNDS_Y[0] < y)
+            & (y < BOUNDS_Y[1])
+            & (BOUNDS_Z[0] < z)
+            & (z < BOUNDS_Z[1])
+        ),
+    )
+    pred_input_dict = {
+        k: v for k, v in samples.items() if k in cfg.MODEL.disp_net.input_keys
+    }
+
+    output_dict = predictor.predict(pred_input_dict, cfg.INFER.batch_size)
+
+    # mapping data to output_keys
+    output_keys = cfg.MODEL.disp_net.output_keys + cfg.MODEL.stress_net.output_keys
+    output_dict = {
+        store_key: output_dict[infer_key]
+        for store_key, infer_key in zip(output_keys, output_dict.keys())
+    }
+    output_dict.update(pred_input_dict)
+
+    vtu.save_vtu_from_dict(
+        osp.join(cfg.output_dir, "vis"),
+        output_dict,
+        cfg.MODEL.disp_net.input_keys,
+        output_keys,
+        1,
+    )
+
+
 @hydra.main(
     version_base=None, config_path="./conf", config_name="forward_analysis.yaml"
 )
@@ -289,8 +364,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/control_arm/inverse_parameter.py b/examples/control_arm/inverse_parameter.py
@@ -265,6 +265,86 @@ def evaluate(cfg: DictConfig):
     solver.visualize()
 
 
+def export(cfg: DictConfig):
+    from paddle.static import InputSpec
+
+    # set model
+    disp_net = ppsci.arch.MLP(**cfg.MODEL.disp_net)
+    stress_net = ppsci.arch.MLP(**cfg.MODEL.stress_net)
+    inverse_lambda_net = ppsci.arch.MLP(**cfg.MODEL.inverse_lambda_net)
+    inverse_mu_net = ppsci.arch.MLP(**cfg.MODEL.inverse_mu_net)
+    # wrap to a model_list
+    model = ppsci.arch.ModelList(
+        (disp_net, stress_net, inverse_lambda_net, inverse_mu_net)
+    )
+
+    # load pretrained model
+    solver = ppsci.solver.Solver(
+        model=model, pretrained_model_path=cfg.INFER.pretrained_model_path
+    )
+
+    # export models
+    input_spec = [
+        {
+            key: InputSpec([None, 1], "float32", name=key)
+            for key in cfg.MODEL.disp_net.input_keys
+        },
+    ]
+    solver.export(input_spec, cfg.INFER.export_path)
+
+
+def inference(cfg: DictConfig):
+    from deploy.python_infer import pinn_predictor
+    from ppsci.visualize import vtu
+
+    # set model predictor
+    predictor = pinn_predictor.PINNPredictor(cfg)
+
+    # set geometry
+    control_arm = ppsci.geometry.Mesh(cfg.GEOM_PATH)
+    # geometry bool operation
+    geo = control_arm
+    geom = {"geo": geo}
+    # set bounds
+    BOUNDS_X, BOUNDS_Y, BOUNDS_Z = control_arm.bounds
+    samples = geom["geo"].sample_interior(
+        cfg.EVAL.batch_size.visualizer_vtu,
+        criteria=lambda x, y, z: (
+            (BOUNDS_X[0] < x)
+            & (x < BOUNDS_X[1])
+            & (BOUNDS_Y[0] < y)
+            & (y < BOUNDS_Y[1])
+            & (BOUNDS_Z[0] < z)
+            & (z < BOUNDS_Z[1])
+        ),
+    )
+    pred_input_dict = {
+        k: v for k, v in samples.items() if k in cfg.MODEL.disp_net.input_keys
+    }
+
+    output_dict = predictor.predict(pred_input_dict, cfg.INFER.batch_size)
+
+    # mapping data to output_keys
+    output_keys = (
+        cfg.MODEL.disp_net.output_keys
+        + cfg.MODEL.stress_net.output_keys
+        + cfg.MODEL.inverse_lambda_net.output_keys
+        + cfg.MODEL.inverse_mu_net.output_keys
+    )
+    output_dict = {
+        store_key: output_dict[infer_key]
+        for store_key, infer_key in zip(output_keys, output_dict.keys())
+    }
+    output_dict.update(pred_input_dict)
+    vtu.save_vtu_from_dict(
+        osp.join(cfg.output_dir, "vis"),
+        output_dict,
+        cfg.MODEL.disp_net.input_keys,
+        output_keys,
+        1,
+    )
+
+
 @hydra.main(
     version_base=None, config_path="./conf", config_name="inverse_parameter.yaml"
 )
@@ -273,8 +353,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__":
