Merge pull request #325 from HydrogenSulfate/okok

support @jit.to_static for euler_beam example

Author: zhiminzhang0830

examples/euler_beam/euler_beam.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import numpy as np
+from paddle import fluid
 
 import ppsci
 from ppsci.autodiff import hessian
@@ -22,6 +22,9 @@
 
 if __name__ == "__main__":
     args = config.parse_args()
+    # enable computation for fourth-order differentiation of matmul
+    fluid.core.set_prim_eager_enabled(True)
+    fluid.core._set_prim_all_enabled(True)
     # set random seed for reproducibility
     ppsci.utils.misc.set_random_seed(42)
     # set training hyper-parameters
@@ -56,49 +59,24 @@
         random="Hammersley",
         name="EQ",
     )
-    bc1 = ppsci.constraint.BoundaryConstraint(
-        {"u0": lambda d: d["u"]},
-        {"u0": 0},
-        geom["interval"],
-        {**dataloader_cfg, "batch_size": 1},
-        ppsci.loss.MSELoss("sum"),
-        criteria=lambda x: np.isclose(x, 0.0),
-        name="BC1",
-    )
-    bc2 = ppsci.constraint.BoundaryConstraint(
-        {"u__x": lambda d: jacobian(d["u"], d["x"])},
-        {"u__x": 0},
-        geom["interval"],
-        {**dataloader_cfg, "batch_size": 1},
-        ppsci.loss.MSELoss("sum"),
-        criteria=lambda x: np.isclose(x, 0.0),
-        name="BC2",
-    )
-    bc3 = ppsci.constraint.BoundaryConstraint(
-        {"u__x__x": lambda d: hessian(d["u"], d["x"])},
-        {"u__x__x": 0},
-        geom["interval"],
-        {**dataloader_cfg, "batch_size": 1},
-        ppsci.loss.MSELoss("sum"),
-        criteria=lambda x: np.isclose(x, 1.0),
-        name="BC3",
-    )
-    bc4 = ppsci.constraint.BoundaryConstraint(
-        {"u__x__x__x": lambda d: jacobian(hessian(d["u"], d["x"]), d["x"])},
-        {"u__x__x__x": 0},
+    bc = ppsci.constraint.BoundaryConstraint(
+        {
+            "u0": lambda d: d["u"][0:1],
+            "u__x": lambda d: jacobian(d["u"], d["x"])[1:2],
+            "u__x__x": lambda d: hessian(d["u"], d["x"])[2:3],
+            "u__x__x__x": lambda d: jacobian(hessian(d["u"], d["x"]), d["x"])[3:4],
+        },
+        {"u0": 0, "u__x": 0, "u__x__x": 0, "u__x__x__x": 0},
         geom["interval"],
-        {**dataloader_cfg, "batch_size": 1},
+        {**dataloader_cfg, "batch_size": 4},
         ppsci.loss.MSELoss("sum"),
-        criteria=lambda x: np.isclose(x, 1.0),
-        name="BC4",
+        evenly=True,
+        name="BC",
     )
     # wrap constraints together
     constraint = {
         pde_constraint.name: pde_constraint,
-        bc1.name: bc1,
-        bc2.name: bc2,
-        bc3.name: bc3,
-        bc4.name: bc4,
+        bc.name: bc,
     }
 
     # set optimizer
@@ -153,6 +131,7 @@ def u_solution_func(out):
         geom=geom,
         validator=validator,
         visualizer=visualizer,
+        to_static=args.to_static,
     )
     # train model
     solver.train()

ppsci/solver/eval.py

@@ -18,7 +18,6 @@
 from paddle import io
 
 from ppsci.solver import printer
-from ppsci.utils import expression
 from ppsci.utils import misc
 from ppsci.utils import profiler
 
@@ -59,17 +58,17 @@ def _eval_by_dataset(solver, epoch_id: int, log_freq: int) -> float:
 
             for v in input_dict.values():
                 v.stop_gradient = False
-            evaluator = expression.ExpressionSolver(
-                _validator.input_keys, _validator.output_keys, solver.model
-            )
-            for output_name, output_formula in _validator.output_expr.items():
-                if output_name in label_dict:
-                    evaluator.add_target_expr(output_formula, output_name)
 
             # forward
             with solver.autocast_context_manager(), solver.no_grad_context_manager():
-                output_dict = evaluator(input_dict)
-                validator_loss = _validator.loss(output_dict, label_dict, weight_dict)
+                output_dict, validator_loss = solver.forward_helper.eval_forward(
+                    _validator.output_expr,
+                    input_dict,
+                    solver.model,
+                    _validator,
+                    label_dict,
+                    weight_dict,
+                )
 
             loss_dict[f"loss({_validator.name})"] = float(validator_loss)
 
@@ -183,16 +182,16 @@ def _eval_by_batch(solver, epoch_id: int, log_freq: int) -> float:
 
             for v in input_dict.values():
                 v.stop_gradient = False
-            evaluator = expression.ExpressionSolver(
-                _validator.input_keys, _validator.output_keys, solver.model
-            )
-            for output_name, output_formula in _validator.output_expr.items():
-                evaluator.add_target_expr(output_formula, output_name)
-
             # forward
             with solver.autocast_context_manager(), solver.no_grad_context_manager():
-                output_dict = evaluator(input_dict)
-                validator_loss = _validator.loss(output_dict, label_dict, weight_dict)
+                output_dict, validator_loss = solver.forward_helper.eval_forward(
+                    _validator.output_expr,
+                    input_dict,
+                    solver.model,
+                    _validator,
+                    label_dict,
+                    weight_dict,
+                )
 
             loss_dict[f"loss({_validator.name})"] = float(validator_loss)
 

ppsci/solver/solver.py

@@ -37,6 +37,7 @@
 
 import ppsci
 from ppsci.utils import config
+from ppsci.utils import expression
 from ppsci.utils import logger
 from ppsci.utils import misc
 from ppsci.utils import save_load
@@ -73,6 +74,7 @@ class Solver:
         compute_metric_by_batch (bool, optional): Whether calculate metrics after each batch during evaluate. Defaults to False.
         eval_with_no_grad (bool, optional): Whether set `stop_gradient=True` for every Tensor if no differentiation
             involved during computation, generally for save GPU memory and accelerate computing. Defaults to False.
+        to_static (bool, optional): Whether enable to_static for forward pass. Defaults to False.
 
     Examples:
         >>> import ppsci
@@ -97,7 +99,7 @@ class Solver:
         ...     "./output",
         ...     opt,
         ...     None,
-        ... )
+        ... )  # doctest: +SKIP
     """
 
     def __init__(
@@ -128,6 +130,7 @@ def __init__(
         checkpoint_path: Optional[str] = None,
         compute_metric_by_batch: bool = False,
         eval_with_no_grad: bool = False,
+        to_static: bool = False,
     ):
         # set model
         self.model = model
@@ -216,6 +219,10 @@ def __init__(
             if isinstance(loaded_metric, dict):
                 self.best_metric.update(loaded_metric)
 
+        # init logger without FileHandler if not initialized before
+        if logger._logger is None:
+            logger.init_logger("ppsci", None)
+
         # choosing an appropriate training function for different optimizers
         if isinstance(self.optimizer, optim.LBFGS):
             self.train_epoch_func = ppsci.solver.train.train_LBFGS_epoch_func
@@ -249,8 +256,13 @@ def __init__(
             if version.Version(paddle.__version__) != version.Version("0.0.0")
             else f"develop({paddle.version.commit[:7]})"
         )
-        if logger._logger is not None:
-            logger.info(f"Using paddlepaddle {paddle_version} on device {self.device}")
+        logger.info(f"Using paddlepaddle {paddle_version} on device {self.device}")
+
+        self.forward_helper = expression.ExpressionSolver()
+
+        # whether enable static for forward pass, default to Fals
+        jit.enable_to_static(to_static)
+        logger.info(f"Set to_static={to_static} for forward computation.")
 
     @staticmethod
     def from_config(cfg: Dict[str, Any]) -> Solver:

ppsci/solver/train.py

@@ -38,36 +38,44 @@ def train_epoch_func(solver, epoch_id: int, log_freq: int):
         reader_cost = 0
         batch_cost = 0
         reader_tic = time.perf_counter()
+
+        input_dicts = []
+        label_dicts = []
+        weight_dicts = []
         for _, _constraint in solver.constraint.items():
             input_dict, label_dict, weight_dict = next(_constraint.data_iter)
-
             # profile code below
             # profiler.add_profiler_step(solver.cfg["profiler_options"])
             if iter_id == 5:
                 # 5 step for warmup
                 for key in solver.train_time_info:
                     solver.train_time_info[key].reset()
             reader_cost += time.perf_counter() - reader_tic
-            total_batch_size += next(iter(input_dict.values())).shape[0]
-
             for v in input_dict.values():
                 v.stop_gradient = False
-            evaluator = expression.ExpressionSolver(
-                _constraint.input_keys, _constraint.output_keys, solver.model
-            )
-            for output_name, output_formula in _constraint.output_expr.items():
-                if output_name in label_dict:
-                    evaluator.add_target_expr(output_formula, output_name)
 
-            # forward for every constraint
-            with solver.autocast_context_manager():
-                output_dict = evaluator(input_dict)
-                constraint_loss = _constraint.loss(output_dict, label_dict, weight_dict)
-                total_loss += constraint_loss
+            # gather each constraint's input, label, weight to a list
+            input_dicts.append(input_dict)
+            label_dicts.append(label_dict)
+            weight_dicts.append(weight_dict)
+            total_batch_size += next(iter(input_dict.values())).shape[0]
+            reader_tic = time.perf_counter()
 
-            loss_dict[_constraint.name] = float(constraint_loss)
+        # forward for every constraint, including model and equation expression
+        with solver.autocast_context_manager():
+            constraint_losses = solver.forward_helper.train_forward(
+                [_constraint.output_expr for _constraint in solver.constraint.values()],
+                input_dicts,
+                solver.model,
+                solver.constraint,
+                label_dicts,
+                weight_dicts,
+            )
 
-            reader_tic = time.perf_counter()
+        # compute loss for each constraint according to its' own output, label and weight
+        for i, _constraint in enumerate(solver.constraint.values()):
+            total_loss += constraint_losses[i]
+            loss_dict[_constraint.name] += float(constraint_losses[i])
 
         if solver.update_freq > 1:
             total_loss = total_loss / solver.update_freq

ppsci/solver/visu.py

@@ -17,7 +17,6 @@
 
 import paddle
 
-from ppsci.utils import expression
 from ppsci.utils import misc
 
 
@@ -52,15 +51,11 @@ def visualize_func(solver, epoch_id: int):
                     batch_input_dict[key] = input_dict[key][st:ed]
                 batch_input_dict[key].stop_gradient = False
 
-            evaluator = expression.ExpressionSolver(
-                _visualizer.input_keys, _visualizer.output_keys, solver.model
-            )
-            for output_key, output_expr in _visualizer.output_expr.items():
-                evaluator.add_target_expr(output_expr, output_key)
-
             # forward
-            with solver.autocast_context_manager(), solver.no_grad_context_manager():
-                batch_output_dict = evaluator(batch_input_dict)
+            with solver.no_grad_context_manager():
+                batch_output_dict = solver.forward_helper.visu_forward(
+                    _visualizer.output_expr, batch_input_dict, solver.model
+                )
 
             # collect batch data
             for key, batch_input in batch_input_dict.items():