Merge branch 'develop' into bubble_net

Author: liaoxin2

docs/zh/examples/volterra_ide.md

@@ -14,6 +14,18 @@
     python volterra_ide.py mode=eval EVAL.pretrained_model_path=https://paddle-org.bj.bcebos.com/paddlescience/models/volterra_ide/volterra_ide_pretrained.pdparams
     ```
 
+=== "模型导出命令"
+
+    ``` sh
+    python volterra_ide.py mode=export
+    ```
+
+=== "模型推理命令"
+
+    ``` sh
+    python volterra_ide.py mode=infer
+    ```
+
 | 预训练模型  | 指标 |
 |:--| :--|
 | [volterra_ide_pretrained.pdparams](https://paddle-org.bj.bcebos.com/paddlescience/models/volterra_ide/volterra_ide_pretrained.pdparams) | loss(L2Rel_Validator): 0.00023 <br> L2Rel.u(L2Rel_Validator): 0.00023 |

examples/ide/conf/volterra_ide.yaml

@@ -61,3 +61,20 @@ EVAL:
   pretrained_model_path: null
   eval_with_no_grad: true
   npoint_eval: 100
+
+INFER:
+  pretrained_model_path: https://paddle-org.bj.bcebos.com/paddlescience/models/volterra_ide/volterra_ide_pretrained.pdparams
+  export_path: ./inference/volterra_ide
+  pdmodel_path: ${INFER.export_path}.pdmodel
+  pdpiparams_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: 64
+  num_cpu_threads: 4
+  batch_size: 16

examples/ide/volterra_ide.py

@@ -254,14 +254,80 @@ def u_solution_func(in_) -> np.ndarray:
     plt.savefig(osp.join(cfg.output_dir, "./Volterra_IDE.png"), dpi=200)
 
 
+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 cfg.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 = {"timedomain": ppsci.geometry.TimeDomain(*cfg.BOUNDS)}
+
+    input_data = geom["timedomain"].uniform_points(cfg.EVAL.npoint_eval)
+    input_dict = {"x": input_data}
+
+    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())
+    }
+
+    def u_solution_func(in_) -> np.ndarray:
+        if isinstance(in_["x"], paddle.Tensor):
+            return paddle.exp(-in_["x"]) * paddle.cosh(in_["x"])
+        return np.exp(-in_["x"]) * np.cosh(in_["x"])
+
+    label_data = u_solution_func({"x": input_data})
+    output_data = output_dict["u"]
+
+    # save result
+    plt.plot(input_data, label_data, "-", label=r"$u(t)$")
+    plt.plot(input_data, output_data, "o", label=r"$\hat{u}(t)$", markersize=4.0)
+    plt.legend()
+    plt.xlabel(r"$t$")
+    plt.ylabel(r"$u$")
+    plt.title(r"$u-t$")
+    plt.savefig("./Volterra_IDE_pred.png", dpi=200)
+
+
 @hydra.main(version_base=None, config_path="./conf", config_name="volterra_ide.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__":

ppsci/data/__init__.py

@@ -82,16 +82,21 @@ def build_dataloader(_dataset, cfg):
         sampler_cfg["batch_size"] = cfg["batch_size"]
         batch_sampler = getattr(io, batch_sampler_cls)(_dataset, **sampler_cfg)
     else:
-        if cfg["batch_size"] != 1:
-            raise ValueError(
-                f"`batch_size` should be 1 when sampler config is None, but got {cfg['batch_size']}."
+        batch_sampler_cls = "BatchSampler"
+        if world_size > 1:
+            batch_sampler_cls = "DistributedBatchSampler"
+            logger.warning(
+                f"Automatically use 'DistributedBatchSampler' instead of "
+                f"'BatchSampler' when world_size({world_size}) > 1."
             )
-        logger.warning(
-            "`batch_size` is set to 1 as neither sampler config nor batch_size is set."
-        )
-        batch_sampler = io.BatchSampler(
+        batch_sampler = getattr(io, batch_sampler_cls)(
             _dataset,
             batch_size=cfg["batch_size"],
+            shuffle=False,
+            drop_last=False,
+        )
+        logger.message(
+            "'shuffle' and 'drop_last' are both set to False in default as sampler config is not specified."
         )
 
     # build collate_fn if specified

ppsci/utils/misc.py

@@ -210,6 +210,11 @@ class Timer(ContextDecorator):
         ...     w = sum(range(0, 10))
         >>> func()  # doctest: +SKIP
 
+        >>> timer = misc.Timer("cost_of_func", auto_print=False)
+        >>> timer.start()
+        >>> func()
+        >>> timer.end()
+        >>> print(f"time cost of 'cost_of_func' is {timer.interval:.2f}")
     """
 
     interval: float  # Time cost for code within Timer context
@@ -220,10 +225,31 @@ def __init__(self, name: str = "Timer", auto_print: bool = True):
         self.auto_print = auto_print
 
     def __enter__(self):
+        paddle.device.synchronize()
         self.start_time = time.perf_counter()
         return self
 
     def __exit__(self, type, value, traceback):
+        paddle.device.synchronize()
+        self.end_time = time.perf_counter()
+        self.interval = self.end_time - self.start_time
+        if self.auto_print:
+            logger.message(f"{self.name}.time_cost = {self.interval:.2f} s")
+
+    def start(self, name: str = "Timer"):
+        """Push a new timer context.
+
+        Args:
+            name (str, optional): Name of code block to be clocked. Defaults to "Timer".
+        """
+        paddle.device.synchronize()
+        self.start_time = time.perf_counter()
+
+    def end(self):
+        """
+        End current timer context and print time cost.
+        """
+        paddle.device.synchronize()
         self.end_time = time.perf_counter()
         self.interval = self.end_time - self.start_time
         if self.auto_print:

ppsci/utils/save_load.py

@@ -146,7 +146,18 @@ def load_checkpoint(
             equation_dict = paddle.load(f"{path}.pdeqn")
 
     # set state dict
-    model.set_state_dict(param_dict)
+    missing_keys, unexpected_keys = model.set_state_dict(param_dict)
+    if missing_keys:
+        logger.warning(
+            f"There are missing keys when loading checkpoint: {missing_keys}, "
+            "and corresponding parameters will be initialized by default."
+        )
+    if unexpected_keys:
+        logger.warning(
+            f"There are redundant keys: {unexpected_keys}, "
+            "and corresponding weights will be ignored."
+        )
+
     optimizer.set_state_dict(optim_dict)
     if grad_scaler is not None:
         grad_scaler.load_state_dict(scaler_dict)

ppsci/utils/symbolic.py

@@ -101,7 +101,7 @@
     sp.sign: paddle.sign,
     sp.ceiling: paddle.ceil,
     sp.floor: paddle.floor,
-    # NOTE: sp.Add and sp.Mul is not included here for un-alignment with sympy
+    # NOTE: sp.Add and sp.Mul is not included here for un-alignment with paddle
     # and are implemented manually in 'OperatorNode._add_operator_func' and
     # 'OperatorNode._mul_operator_func'
 }
@@ -711,15 +711,15 @@ def lambdify(
             such as 'momentum_x'. Defaults to None.
         create_graph (bool, optional): Whether to create the gradient graphs of
             the computing process. When it is True, higher order derivatives are
-            supported to compute; when it is False, the gradient graphs of the
+            supported to compute. When it is False, the gradient graphs of the
             computing process would be discarded. Defaults to True.
         retain_graph (Optional[bool]): Whether to retain the forward graph which
             is used to calculate the gradient. When it is True, the graph would
             be retained, in which way users can calculate backward twice for the
             same graph. When it is False, the graph would be freed. Defaults to None,
             which means it is equal to `create_graph`.
         fuse_derivative (bool, optional): Whether to fuse the derivative nodes.
-            for example, if `expr` is 'Derivative(u, x) + Derivative(u, y)'
+            For example, if `expr` is 'Derivative(u, x) + Derivative(u, y)'
             It will compute grad(u, x) + grad(u, y) if fuse_derivative=False,
             else will compute sum(grad(u, [x, y])) if fuse_derivative=True as is more
             efficient in backward-graph. Defaults to False, as it is experimental so not