refine expression code and support train/eval/visu for euler_beam in jit.to_static

Author: HydrogenSulfate

examples/darcy/darcy2d.py

@@ -13,15 +13,12 @@
 # limitations under the License.
 
 import numpy as np
-from paddle import fluid
 
 import ppsci
 from ppsci.utils import config
 from ppsci.utils import logger
 
 if __name__ == "__main__":
-    fluid.core._set_prim_all_enabled(True)
-
     args = config.parse_args()
     # set random seed for reproducibility
     ppsci.utils.misc.set_random_seed(42)

examples/euler_beam/euler_beam.py

@@ -126,11 +126,12 @@ def u_solution_func(out):
         epochs=EPOCHS,
         iters_per_epoch=ITERS_PER_EPOCH,
         eval_during_train=True,
-        eval_freq=1000,
+        eval_freq=10,
         equation=equation,
         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
 
@@ -62,10 +61,14 @@ def _eval_by_dataset(solver, epoch_id: int, log_freq: int) -> float:
 
             # forward
             with solver.autocast_context_manager(), solver.no_grad_context_manager():
-                output_dict = solver.expr_helper(
-                    _validator.output_expr, input_dict, solver.model
+                output_dict, validator_loss = solver.forward_helper.eval_forward(
+                    _validator.output_expr,
+                    input_dict,
+                    solver.model,
+                    _validator,
+                    label_dict,
+                    weight_dict,
                 )
-                validator_loss = _validator.loss(output_dict, label_dict, weight_dict)
 
             loss_dict[f"loss({_validator.name})"] = float(validator_loss)
 
@@ -179,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

@@ -74,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
@@ -129,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
@@ -253,7 +255,12 @@ def __init__(
         if logger._logger is not None:
             logger.info(f"Using paddlepaddle {paddle_version} on device {self.device}")
 
-        self.expr_helper = expression.ExpressionSolver()
+        self.forward_helper = expression.ExpressionSolver()
+
+        # whether enable static for forward pass, default to Fals
+        if to_static:
+            jit.enable_to_static(to_static)
+            logger.info("Enable to_static for forward computation.")
 
     @staticmethod
     def from_config(cfg: Dict[str, Any]) -> Solver:

ppsci/solver/train.py

@@ -39,9 +39,9 @@ def train_epoch_func(solver, epoch_id: int, log_freq: int):
         batch_cost = 0
         reader_tic = time.perf_counter()
 
-        input_dict_list = []
-        label_dict_list = []
-        weight_dict_list = []
+        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
@@ -51,35 +51,31 @@ def train_epoch_func(solver, epoch_id: int, log_freq: int):
                 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
 
             # gather each constraint's input, label, weight to a list
-            input_dict_list.append(input_dict)
-            label_dict_list.append(label_dict)
-            weight_dict_list.append(weight_dict)
-
+            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()
 
-        for x in input_dict_list:
-            for v in x.values():
-                v.stop_gradient = False
-
         # forward for every constraint, including model and equation expression
         with solver.autocast_context_manager():
-            constraint_losses = solver.expr_helper(
+            constraint_losses = solver.forward_helper.train_forward(
                 [_constraint.output_expr for _constraint in solver.constraint.values()],
-                input_dict_list,
+                input_dicts,
                 solver.model,
                 solver.constraint,
-                label_dict_list,
-                weight_dict_list,
+                label_dicts,
+                weight_dicts,
             )
 
         # compute loss for each constraint according to its' own output, label and weight
-        for i, (_, _constraint) in enumerate(solver.constraint.items()):
-            constraint_loss = constraint_losses[i]
-            total_loss += constraint_loss
-            loss_dict[_constraint.name] += float(constraint_loss)
+        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
 
 
@@ -54,7 +53,7 @@ def visualize_func(solver, epoch_id: int):
 
             # forward
             with solver.no_grad_context_manager():
-                batch_output_dict = solver.expr_helper(
+                batch_output_dict = solver.forward_helper.visu_forward(
                     _visualizer.output_expr, batch_input_dict, solver.model
                 )
 

ppsci/utils/config.py

@@ -179,6 +179,11 @@ def parse_args():
     parser = argparse.ArgumentParser("paddlescience running script")
     parser.add_argument("-e", "--epochs", type=int, help="training epochs")
     parser.add_argument("-o", "--output_dir", type=str, help="output directory")
+    parser.add_argument(
+        "--to_static",
+        action="store_true",
+        help="whether enable to_static for forward computation",
+    )
 
     args = parser.parse_args()
     return args

ppsci/utils/expression.py

@@ -12,64 +12,128 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import TYPE_CHECKING
+from typing import Callable
+from typing import Dict
+from typing import Tuple
+
+import paddle
 from paddle import jit
 from paddle import nn
 
+if TYPE_CHECKING:
+    from ppsci import constraint
+    from ppsci import validate
+
 from ppsci.autodiff import clear
 
 
 class ExpressionSolver(nn.Layer):
     """Expression Solver
 
-    Args:
-        input_keys (Dict[str]):Names of input keys.
-        output_keys (Dict[str]):Names of output keys.
-        model (nn.Layer): Model to get output variables from input variables.
-
     Examples:
         >>> import ppsci
         >>> model = ppsci.arch.MLP(("x", "y"), ("u", "v"), 5, 128)
-        >>> expr_solver = ExpressionSolver(("x", "y"), ("u", "v"), model)
+        >>> expr_solver = ExpressionSolver()
     """
 
     def __init__(self):
         super().__init__()
 
     @jit.to_static
-    def forward(
+    def train_forward(
         self,
-        expr_dict_list,
-        input_dict_list,
-        model,
-        constraint,
-        label_dict_list,
-        weight_dict_list,
+        expr_dicts: Tuple[Dict[str, Callable], ...],
+        input_dicts: Tuple[Dict[str, paddle.Tensor], ...],
+        model: nn.Layer,
+        constraint: Dict[str, "constraint.Constraint"],
+        label_dicts: Tuple[Dict[str, paddle.Tensor], ...],
+        weight_dicts: Tuple[Dict[str, paddle.Tensor], ...],
     ):
-        output_dict_list = []
-        for i, expr_dict in enumerate(expr_dict_list):
+        output_dicts = []
+        for i, expr_dict in enumerate(expr_dicts):
             # model forward
             if callable(next(iter(expr_dict.values()))):
-                output_dict = model(input_dict_list[i])
+                output_dict = model(input_dicts[i])
 
             # equation forward
             for name, expr in expr_dict.items():
+                if name not in label_dicts[i]:
+                    continue
                 if callable(expr):
-                    output_dict[name] = expr({**output_dict, **input_dict_list[i]})
+                    output_dict[name] = expr({**output_dict, **input_dicts[i]})
                 else:
                     raise TypeError(f"expr type({type(expr)}) is invalid")
 
-            output_dict_list.append(output_dict)
+            output_dicts.append(output_dict)
 
             # clear differentiation cache
             clear()
 
         # compute loss for each constraint according to its' own output, label and weight
         constraint_losses = []
-        for i, (_, _constraint) in enumerate(constraint.items()):
+        for i, _constraint in enumerate(constraint.values()):
             constraint_loss = _constraint.loss(
-                output_dict_list[i],
-                label_dict_list[i],
-                weight_dict_list[i],
+                output_dicts[i],
+                label_dicts[i],
+                weight_dicts[i],
             )
             constraint_losses.append(constraint_loss)
         return constraint_losses
+
+    @jit.to_static
+    def eval_forward(
+        self,
+        expr_dict: Dict[str, Callable],
+        input_dict: Dict[str, Callable],
+        model: nn.Layer,
+        validator: "validate.Validator",
+        label_dict: Dict[str, Callable],
+        weight_dict: Dict[str, Callable],
+    ):
+        # model forward
+        if callable(next(iter(expr_dict.values()))):
+            output_dict = model(input_dict)
+
+        # equation forward
+        for name, expr in expr_dict.items():
+            if name not in label_dict:
+                continue
+            if callable(expr):
+                output_dict[name] = expr({**output_dict, **input_dict})
+            else:
+                raise TypeError(f"expr type({type(expr)}) is invalid")
+
+        # clear differentiation cache
+        clear()
+
+        # compute loss for each validator according to its' own output, label and weight
+        validator_loss = validator.loss(
+            output_dict,
+            label_dict,
+            weight_dict,
+        )
+        return output_dict, validator_loss
+
+    def visu_forward(
+        self,
+        expr_dict: Dict[str, Callable],
+        input_dict: Dict[str, Callable],
+        model: nn.Layer,
+    ):
+        # model forward
+        if callable(next(iter(expr_dict.values()))):
+            output_dict = model(input_dict)
+
+        # equation forward
+        for name, expr in expr_dict.items():
+            if callable(expr):
+                output_dict[name] = expr({**output_dict, **input_dict})
+            else:
+                raise TypeError(f"expr type({type(expr)}) is invalid")
+
+        # clear differentiation cache
+        clear()
+
+        # compute loss for each validator according to its' own output, label and weight
+        return output_dict

ppsci/utils/misc.py

@@ -255,10 +255,12 @@ def function_with_eval_state(self, *args, **kwargs):
             self.model.eval()
 
         # run func in eval mode
-        func(self, *args, **kwargs)
+        result = func(self, *args, **kwargs)
 
         # restore state
         if train_state:
             self.model.train()
 
+        return result
+
     return function_with_eval_state