update runable code(wrap subloss computation in static

Author: HydrogenSulfate

ppsci/solver/train.py

@@ -66,19 +66,18 @@ def train_epoch_func(solver, epoch_id: int, log_freq: int):
 
         # forward for every constraint, including model and equation expression
         with solver.autocast_context_manager():
-            output_dict_list = solver.expr_helper(
+            constraint_losses = solver.expr_helper(
                 [_constraint.output_expr for _constraint in solver.constraint.values()],
                 input_dict_list,
                 solver.model,
+                solver.constraint,
+                label_dict_list,
+                weight_dict_list,
             )
 
         # compute loss for each constraint according to its' own output, label and weight
         for i, (_, _constraint) in enumerate(solver.constraint.items()):
-            constraint_loss = _constraint.loss(
-                output_dict_list[i],
-                label_dict_list[i],
-                weight_dict_list[i],
-            )
+            constraint_loss = constraint_losses[i]
             total_loss += constraint_loss
             loss_dict[_constraint.name] += float(constraint_loss)
 

ppsci/utils/expression.py

@@ -36,7 +36,15 @@ def __init__(self):
         super().__init__()
 
     @jit.to_static
-    def forward(self, expr_dict_list, input_dict_list, model):
+    def forward(
+        self,
+        expr_dict_list,
+        input_dict_list,
+        model,
+        constraint,
+        label_dict_list,
+        weight_dict_list,
+    ):
         output_dict_list = []
         for i, expr_dict in enumerate(expr_dict_list):
             # model forward
@@ -55,4 +63,13 @@ def forward(self, expr_dict_list, input_dict_list, model):
             # clear differentiation cache
             clear()
 
-        return output_dict_list
+        # compute loss for each constraint according to its' own output, label and weight
+        constraint_losses = []
+        for i, (_, _constraint) in enumerate(constraint.items()):
+            constraint_loss = _constraint.loss(
+                output_dict_list[i],
+                label_dict_list[i],
+                weight_dict_list[i],
+            )
+            constraint_losses.append(constraint_loss)
+        return constraint_losses