[Refine] Refine loss and metric module (#919)

* return loss dict instead of loss summation for all loss.forward

* adapt all mtl module for Dict[str, Tensor] type of input losses

* fix

* remove 'area' in Constriant.output_keys

* fix eval.py

* fix code

* fix examples in func.py

* fix examples in func.py

* Fix for MSELossWithL2Decay and train_enn.py

* fix doctest in loss/mse.py

* fix epnn

* fix

Author: HydrogenSulfate

examples/RegAE/RegAE.py

@@ -56,7 +56,10 @@ def loss_expr(output_dict, label_dict, weight_dict=None):
         base = paddle.exp(2.0 * log_sigma) + paddle.pow(mu, 2) - 1.0 - 2.0 * log_sigma
         KLLoss = 0.5 * paddle.sum(base) / mu.shape[0]
 
-        return F.mse_loss(output_dict["decoder_z"], label_dict["p_train"]) + KLLoss
+        return {
+            "decode_loss": F.mse_loss(output_dict["decoder_z"], label_dict["p_train"])
+            + KLLoss
+        }
 
     # set constraint
     sup_constraint = ppsci.constraint.SupervisedConstraint(

examples/amgnet/amgnet_airfoil.py

@@ -37,7 +37,7 @@ def train_mse_func(
     label_dict: Dict[str, "pgl.Graph"],
     *args,
 ) -> paddle.Tensor:
-    return F.mse_loss(output_dict["pred"], label_dict["label"].y)
+    return {"pred": F.mse_loss(output_dict["pred"], label_dict["label"].y)}
 
 
 def eval_rmse_func(

examples/amgnet/amgnet_cylinder.py

@@ -37,7 +37,7 @@ def train_mse_func(
     label_dict: Dict[str, "pgl.Graph"],
     *args,
 ) -> paddle.Tensor:
-    return F.mse_loss(output_dict["pred"], label_dict["label"].y)
+    return {"pred": F.mse_loss(output_dict["pred"], label_dict["label"].y)}
 
 
 def eval_rmse_func(

examples/cfdgcn/cfdgcn.py

@@ -33,7 +33,7 @@ def train_mse_func(
     label_dict: Dict[str, "pgl.Graph"],
     *args,
 ) -> paddle.Tensor:
-    return F.mse_loss(output_dict["pred"], label_dict["label"].y)
+    return {"pred": F.mse_loss(output_dict["pred"], label_dict["label"].y)}
 
 
 def eval_rmse_func(

examples/deepcfd/deepcfd.py

@@ -243,7 +243,7 @@ def loss_expr(
         loss_v = (output[:, 1:2, :, :] - y[:, 1:2, :, :]) ** 2
         loss_p = (output[:, 2:3, :, :] - y[:, 2:3, :, :]).abs()
         loss = (loss_u + loss_v + loss_p) / CHANNELS_WEIGHTS
-        return loss.sum()
+        return {"output": loss.sum()}
 
     sup_constraint = ppsci.constraint.SupervisedConstraint(
         {

examples/deephpms/burgers.py

@@ -31,7 +31,7 @@
 
 def pde_loss_func(output_dict, *args):
     losses = F.mse_loss(output_dict["f_pde"], output_dict["du_t"], "sum")
-    return losses
+    return {"pde": losses}
 
 
 def pde_l2_rel_func(output_dict, *args):
@@ -53,7 +53,7 @@ def boundary_loss_func(output_dict, *args):
 
     losses = F.mse_loss(u_lb, u_ub, "sum")
     losses += F.mse_loss(du_x_lb, du_x_ub, "sum")
-    return losses
+    return {"boundary": losses}
 
 
 def train(cfg: DictConfig):

examples/deephpms/korteweg_de_vries.py

@@ -31,7 +31,7 @@
 
 def pde_loss_func(output_dict, *args):
     losses = F.mse_loss(output_dict["f_pde"], output_dict["du_t"], "sum")
-    return losses
+    return {"pde": losses}
 
 
 def pde_l2_rel_func(output_dict, *args):
@@ -56,7 +56,7 @@ def boundary_loss_func(output_dict, *args):
     losses = F.mse_loss(u_lb, u_ub, "sum")
     losses += F.mse_loss(du_x_lb, du_x_ub, "sum")
     losses += F.mse_loss(du_xx_lb, du_xx_ub, "sum")
-    return losses
+    return {"boundary": losses}
 
 
 def train(cfg: DictConfig):

examples/deephpms/kuramoto_sivashinsky.py

@@ -31,7 +31,7 @@
 
 def pde_loss_func(output_dict, *args):
     losses = F.mse_loss(output_dict["f_pde"], output_dict["du_t"], "sum")
-    return losses
+    return {"pde": losses}
 
 
 def pde_l2_rel_func(output_dict, *args):
@@ -59,7 +59,7 @@ def boundary_loss_func(output_dict, *args):
     losses += F.mse_loss(du_x_lb, du_x_ub, "sum")
     losses += F.mse_loss(du_xx_lb, du_xx_ub, "sum")
     losses += F.mse_loss(du_xxx_lb, du_xxx_ub, "sum")
-    return losses
+    return {"boundary": losses}
 
 
 def train(cfg: DictConfig):

examples/deephpms/navier_stokes.py

@@ -31,7 +31,7 @@
 
 def pde_loss_func(output_dict, *args):
     losses = F.mse_loss(output_dict["f_pde"], output_dict["dw_t"], "sum")
-    return losses
+    return {"pde": losses}
 
 
 def pde_l2_rel_func(output_dict, *args):

examples/deephpms/schrodinger.py

@@ -32,7 +32,7 @@
 def pde_loss_func(output_dict, *args):
     losses = F.mse_loss(output_dict["f_pde"], output_dict["du_t"], "sum")
     losses += F.mse_loss(output_dict["g_pde"], output_dict["dv_t"], "sum")
-    return losses
+    return {"pde": losses}
 
 
 def pde_l2_rel_func(output_dict, *args):
@@ -62,7 +62,7 @@ def boundary_loss_func(output_dict, *args):
     losses += F.mse_loss(v_lb, v_ub, "sum")
     losses += F.mse_loss(du_x_lb, du_x_ub, "sum")
     losses += F.mse_loss(dv_x_lb, dv_x_ub, "sum")
-    return losses
+    return {"boundary": losses}
 
 
 def sol_l2_rel_func(output_dict, label_dict):