【PPSCI Export&Infer No.11-12】 (#883)

wufei2 · HydrogenSulfate · web-flow · commit 1af86e17d3a1 · 2024-05-10T20:49:51.000+08:00
* fix doc bugs

* fix codestyle bugs

* 【PPSCI Export&amp;Infer No.15-16】

* fix codestyle bug for PPSCI Export&amp;Infer No.15-16】

* fix codestyle bugs for 【PPSCI Export&amp;Infer No.15-16】

* fix codestyle bugs for 【PPSCI Export&amp;Infer No.15-16】

* fix codestyle bugs for 【PPSCI Export&amp;Infer No.15-16】

* fix bugs for 【PPSCI Export&amp;Infer No.15-16】

* fix codestyle bugs

* 【PPSCI Export&amp;Infer No.11-12】

* change predictor

* fix bugs in change predictor

* cancel extra doc commit

* fix codestyle bugs

* Update examples/cylinder/2d_unsteady/cylinder2d_unsteady_Re100.py

* Update examples/cylinder/2d_unsteady/transformer_physx/train_transformer.py

* Update examples/cylinder/2d_unsteady/cylinder2d_unsteady_Re100.py

* Update examples/cylinder/2d_unsteady/transformer_physx/train_transformer.py

* Update examples/cylinder/2d_unsteady/transformer_physx/train_transformer.py

* cancel extra changes

* cancel extra changes

* update examples/cylinder/2d_unsteady/transformer_physx/conf/transformer.yaml

---------

Co-authored-by: HydrogenSulfate &lt;490868991@qq.com&gt;
diff --git a/docs/zh/examples/cylinder2d_unsteady.md b/docs/zh/examples/cylinder2d_unsteady.md
@@ -26,6 +26,24 @@
     python cylinder2d_unsteady_Re100.py mode=eval EVAL.pretrained_model_path=https://paddle-org.bj.bcebos.com/paddlescience/models/cylinder2d_unsteady_Re100/cylinder2d_unsteady_Re100_pretrained.pdparams
     ```
 
+    === "模型导出命令"
+
+    ``` sh
+    python cylinder2d_unsteady_Re100.py mode=export
+    ```
+
+=== "模型推理命令"
+
+    ``` sh
+    # linux
+    wget -nc https://paddle-org.bj.bcebos.com/paddlescience/datasets/cylinder2d_unsteady_Re100/cylinder2d_unsteady_Re100_dataset.tar
+    # windows
+    # curl https://paddle-org.bj.bcebos.com/paddlescience/datasets/cylinder2d_unsteady_Re100/cylinder2d_unsteady_Re100_dataset.tar --output cylinder2d_unsteady_Re100_dataset.tar
+    # unzip it
+    tar -xvf cylinder2d_unsteady_Re100_dataset.tar
+    python cylinder2d_unsteady_Re100.py mode=infer
+    ```
+
 | 预训练模型  | 指标 |
 |:--| :--|
 | [cylinder2d_unsteady_Re100_pretrained.pdparams](https://paddle-org.bj.bcebos.com/paddlescience/models/cylinder2d_unsteady_Re100/cylinder2d_unsteady_Re100_pretrained.pdparams) | loss(Residual): 0.00398<br>MSE.continuity(Residual): 0.00126<br>MSE.momentum_x(Residual): 0.00151<br>MSE.momentum_y(Residual): 0.00120 |
diff --git a/docs/zh/examples/cylinder2d_unsteady_transformer_physx.md b/docs/zh/examples/cylinder2d_unsteady_transformer_physx.md
@@ -27,6 +27,23 @@
     python train_enn.py mode=eval EVAL.pretrained_model_path=https://paddle-org.bj.bcebos.com/paddlescience/models/cylinder/cylinder_pretrained.pdparams
     python train_transformer.py mode=eval EVAL.pretrained_model_path=https://paddle-org.bj.bcebos.com/paddlescience/models/cylinder/cylinder_transformer_pretrained.pdparams EMBEDDING_MODEL_PATH=https://paddle-org.bj.bcebos.com/paddlescience/models/cylinder/cylinder_pretrained.pdparams
     ```
+    === "模型导出命令"
+
+    ``` sh
+    python train_transformer.py mode=export EMBEDDING_MODEL_PATH=https://paddle-org.bj.bcebos.com/paddlescience/models/cylinder/cylinder_pretrained.pdparams
+    ```
+
+=== "模型推理命令"
+
+    ``` sh
+    # linux
+    wget -nc https://paddle-org.bj.bcebos.com/paddlescience/datasets/transformer_physx/cylinder_training.hdf5 -P ./datasets/
+    wget -nc https://paddle-org.bj.bcebos.com/paddlescience/datasets/transformer_physx/cylinder_valid.hdf5 -P ./datasets/
+    # windows
+    # curl https://paddle-org.bj.bcebos.com/paddlescience/datasets/transformer_physx/cylinder_training.hdf5 --output ./datasets/cylinder_training.hdf5
+    # curl https://paddle-org.bj.bcebos.com/paddlescience/datasets/transformer_physx/cylinder_valid.hdf5 --output ./datasets/cylinder_valid.hdf5
+    python train_transformer.py mode=infer
+    ```
 
 | 模型 | MSE |
 | :-- | :-- |
diff --git a/examples/cylinder/2d_unsteady/conf/cylinder2d_unsteady_Re100.yaml b/examples/cylinder/2d_unsteady/conf/cylinder2d_unsteady_Re100.yaml
@@ -70,3 +70,21 @@ TRAIN:
 EVAL:
   batch_size: 10240
   pretrained_model_path: null
+
+# inference settings
+INFER:
+  pretrained_model_path: https://paddle-org.bj.bcebos.com/paddlescience/models/cylinder2d_unsteady_Re100/cylinder2d_unsteady_Re100_pretrained.pdparams
+  export_path: ./inference/cylinder2d_unsteady_Re100
+  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: 10240
+  num_cpu_threads: 10
+  batch_size: 10240
diff --git a/examples/cylinder/2d_unsteady/cylinder2d_unsteady_Re100.py b/examples/cylinder/2d_unsteady/cylinder2d_unsteady_Re100.py
@@ -311,6 +311,59 @@ def evaluate(cfg: DictConfig):
     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 time-geometry
+    geom = {
+        "time_rect_eval": ppsci.geometry.PointCloud(
+            reader.load_csv_file(
+                cfg.DOMAIN_EVAL_PATH,
+                ("t", "x", "y"),
+            ),
+            ("t", "x", "y"),
+        ),
+    }
+    NPOINT_EVAL = (
+        cfg.NPOINT_PDE + cfg.NPOINT_INLET_CYLINDER + cfg.NPOINT_OUTLET
+    ) * cfg.NUM_TIMESTAMPS
+    input_dict = geom["time_rect_eval"].sample_interior(NPOINT_EVAL, evenly=True)
+    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.MODEL.output_keys, output_dict.keys())
+    }
+
+    ppsci.visualize.save_vtu_from_dict(
+        "./cylinder2d_unsteady_Re100_pred.vtu",
+        {**input_dict, **output_dict},
+        input_dict.keys(),
+        cfg.MODEL.output_keys,
+        cfg.NUM_TIMESTAMPS,
+    )
+
+
 @hydra.main(
     version_base=None,
     config_path="./conf",
@@ -321,8 +374,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/cylinder/2d_unsteady/transformer_physx/conf/transformer.yaml b/examples/cylinder/2d_unsteady/transformer_physx/conf/transformer.yaml
@@ -28,6 +28,7 @@ TRAIN_BLOCK_SIZE: 16
 VALID_BLOCK_SIZE: 256
 TRAIN_FILE_PATH: ./datasets/cylinder_training.hdf5
 VALID_FILE_PATH: ./datasets/cylinder_valid.hdf5
+log_freq: 20
 
 # set working condition
 EMBEDDING_MODEL_PATH: ./outputs_cylinder2d_unsteady_transformer_physx_enn/checkpoints/latest
@@ -37,7 +38,7 @@ VIS_DATA_NUMS: 1
 MODEL:
   input_keys: ["embeds"]
   output_keys: ["pred_embeds"]
-  num_layers: 4
+  num_layers: 6
   num_ctx: 16
   embed_size: 128
   num_heads: 4
@@ -63,3 +64,20 @@ TRAIN:
 EVAL:
   batch_size: 16
   pretrained_model_path: null
+
+INFER:
+  pretrained_model_path: https://paddle-org.bj.bcebos.com/paddlescience/models/cylinder/cylinder_transformer_pretrained.pdparams
+  export_path: ./inference/cylinder_transformer
+  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: false
+  min_subgraph_size: 10
+  gpu_mem: 4000
+  gpu_id: 0
+  max_batch_size: 64
+  num_cpu_threads: 4
+  batch_size: 16
diff --git a/examples/cylinder/2d_unsteady/transformer_physx/train_transformer.py b/examples/cylinder/2d_unsteady/transformer_physx/train_transformer.py
@@ -262,14 +262,91 @@ def evaluate(cfg: DictConfig):
     solver.visualize()
 
 
+def export(cfg: DictConfig):
+    # set model
+    embedding_model = build_embedding_model(cfg.EMBEDDING_MODEL_PATH)
+    model_cfg = {
+        **cfg.MODEL,
+        "embedding_model": embedding_model,
+        "input_keys": ["states"],
+        "output_keys": ["pred_states"],
+    }
+    model = ppsci.arch.PhysformerGPT2(**model_cfg)
+
+    # 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, 16, 128], "float32", name=key)
+            for key in model.input_keys
+        },
+    ]
+
+    solver.export(input_spec, cfg.INFER.export_path)
+
+
+def inference(cfg: DictConfig):
+    from deploy import python_infer
+
+    predictor = python_infer.GeneralPredictor(cfg)
+
+    dataset_cfg = {
+        "name": "CylinderDataset",
+        "file_path": cfg.VALID_FILE_PATH,
+        "input_keys": cfg.MODEL.input_keys,
+        "label_keys": cfg.MODEL.output_keys,
+        "block_size": cfg.VALID_BLOCK_SIZE,
+        "stride": 1024,
+    }
+
+    dataset = ppsci.data.dataset.build_dataset(dataset_cfg)
+
+    input_dict = {
+        "states": dataset.data[: cfg.VIS_DATA_NUMS, :-1],
+    }
+
+    output_dict = predictor.predict(input_dict)
+
+    # mapping data to cfg.INFER.output_keys
+    output_keys = ["pred_states"]
+    output_dict = {
+        store_key: output_dict[infer_key]
+        for store_key, infer_key in zip(output_keys, output_dict.keys())
+    }
+
+    input_dict = {
+        "states": dataset.data[: cfg.VIS_DATA_NUMS, 1:],
+    }
+
+    data_dict = {**input_dict, **output_dict}
+    for i in range(cfg.VIS_DATA_NUMS):
+        ppsci.visualize.save_plot_from_3d_dict(
+            f"./cylinder_transformer_pred_{i}",
+            {key: value[i] for key, value in data_dict.items()},
+            ("states", "pred_states"),
+        )
+
+
 @hydra.main(version_base=None, config_path="./conf", config_name="transformer.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/ppsci/arch/physx_transformer.py b/ppsci/arch/physx_transformer.py
@@ -28,6 +28,7 @@
 from paddle.nn.initializer import Normal
 
 from ppsci.arch import base
+from ppsci.arch.embedding_koopman import CylinderEmbedding
 
 zeros_ = Constant(value=0.0)
 ones_ = Constant(value=1.0)
@@ -387,7 +388,10 @@ def forward(self, x):
             x = self._input_transform(x)
         x_tensor = self.concat_to_tensor(x, self.input_keys, axis=-1)
         if self.embedding_model is not None:
-            x_tensor = self.embedding_model.encoder(x_tensor)
+            if isinstance(self.embedding_model, CylinderEmbedding):
+                x_tensor = self.embedding_model.encoder(x_tensor, x["visc"])
+            else:
+                x_tensor = self.embedding_model.encoder(x_tensor)
 
         if self.training:
             y = self.forward_tensor(x_tensor)
