avoid multiplication with weight 1 if weight is None for reducing com…

…putation overhead (PaddlePaddle#714)

ppsci/constraint/boundary_constraint.py

@@ -131,8 +131,9 @@ def __init__(
                 raise NotImplementedError(f"type of {type(value)} is invalid yet.")
 
         # prepare weight
-        weight = {key: np.ones_like(next(iter(label.values()))) for key in label}
+        weight = None
         if weight_dict is not None:
+            weight = {key: np.ones_like(next(iter(label.values()))) for key in label}
             for key, value in weight_dict.items():
                 if isinstance(value, (int, float)):
                     weight[key] = np.full_like(next(iter(label.values())), value)

ppsci/constraint/initial_constraint.py

@@ -138,8 +138,9 @@ def __init__(
                 raise NotImplementedError(f"type of {type(value)} is invalid yet.")
 
         # prepare weight
-        weight = {key: np.ones_like(next(iter(label.values()))) for key in label}
+        weight = None
         if weight_dict is not None:
+            weight = {key: np.ones_like(next(iter(label.values()))) for key in label}
             for key, value in weight_dict.items():
                 if isinstance(value, (int, float)):
                     weight[key] = np.full_like(next(iter(label.values())), value)

ppsci/constraint/integral_constraint.py

@@ -144,8 +144,9 @@ def __init__(
 
         # prepare weight
         # shape of each weight is [batch_size, ndim]
-        weight = {key: np.ones_like(next(iter(label.values()))) for key in label}
+        weight = None
         if weight_dict is not None:
+            weight = {key: np.ones_like(next(iter(label.values()))) for key in label}
             for key, value in weight_dict.items():
                 if isinstance(value, (int, float)):
                     weight[key] = np.full_like(next(iter(label.values())), value)

ppsci/constraint/interior_constraint.py

@@ -135,8 +135,9 @@ def __init__(
                 raise NotImplementedError(f"type of {type(value)} is invalid yet.")
 
         # prepare weight
-        weight = {key: np.ones_like(next(iter(label.values()))) for key in label}
+        weight = None
         if weight_dict is not None:
+            weight = {key: np.ones_like(next(iter(label.values()))) for key in label}
             for key, value in weight_dict.items():
                 if isinstance(value, str):
                     if value == "sdf":

ppsci/constraint/periodic_constraint.py

@@ -137,8 +137,9 @@ def __init__(
             )
 
         # # prepare weight, keep weight the same shape as input_periodic
-        weight = {key: np.ones_like(next(iter(label.values()))) for key in label}
+        weight = None
         if weight_dict is not None:
+            weight = {key: np.ones_like(next(iter(label.values()))) for key in label}
             for key, value in weight_dict.items():
                 if isinstance(value, (int, float)):
                     weight[key] = np.full_like(next(iter(label.values())), value)

ppsci/data/dataset/csv_dataset.py

@@ -111,9 +111,11 @@ def __init__(
         }
 
         # prepare weights
-        self.weight = {
-            key: np.ones_like(next(iter(self.label.values()))) for key in self.label
-        }
+        self.weight = (
+            {key: np.ones_like(next(iter(self.label.values()))) for key in self.label}
+            if weight_dict is not None
+            else {}
+        )
         if weight_dict is not None:
             for key, value in weight_dict.items():
                 if isinstance(value, (int, float)):
@@ -231,9 +233,11 @@ def __init__(
         }
 
         # prepare weights
-        self.weight = {
-            key: np.ones_like(next(iter(self.label.values()))) for key in self.label
-        }
+        self.weight = (
+            {key: np.ones_like(next(iter(self.label.values()))) for key in self.label}
+            if weight_dict is not None
+            else {}
+        )
         if weight_dict is not None:
             for key, value in weight_dict.items():
                 if isinstance(value, (int, float)):

ppsci/data/dataset/mat_dataset.py

@@ -111,9 +111,11 @@ def __init__(
         }
 
         # prepare weights
-        self.weight = {
-            key: np.ones_like(next(iter(self.label.values()))) for key in self.label
-        }
+        self.weight = (
+            {key: np.ones_like(next(iter(self.label.values()))) for key in self.label}
+            if weight_dict is not None
+            else {}
+        )
         if weight_dict is not None:
             for key, value in weight_dict.items():
                 if isinstance(value, (int, float)):
@@ -231,9 +233,11 @@ def __init__(
         }
 
         # prepare weights
-        self.weight = {
-            key: np.ones_like(next(iter(self.label.values()))) for key in self.label
-        }
+        self.weight = (
+            {key: np.ones_like(next(iter(self.label.values()))) for key in self.label}
+            if weight_dict is not None
+            else {}
+        )
         if weight_dict is not None:
             for key, value in weight_dict.items():
                 if isinstance(value, (int, float)):

ppsci/loss/integral.py

@@ -85,7 +85,7 @@ def forward(self, output_dict, label_dict, weight_dict=None):
                 label_dict[key],
                 "none",
             )
-            if weight_dict:
+            if weight_dict and key in weight_dict:
                 loss *= weight_dict[key]
 
             if self.reduction == "sum":

ppsci/loss/l1.py

@@ -87,7 +87,7 @@ def forward(self, output_dict, label_dict, weight_dict=None):
         losses = 0.0
         for key in label_dict:
             loss = F.l1_loss(output_dict[key], label_dict[key], "none")
-            if weight_dict:
+            if weight_dict and key in weight_dict:
                 loss *= weight_dict[key]
 
             if "area" in output_dict:
@@ -181,7 +181,7 @@ def forward(self, output_dict, label_dict, weight_dict=None):
             loss = F.l1_loss(
                 output_dict[key][:n_output], output_dict[key][n_output:], "none"
             )
-            if weight_dict:
+            if weight_dict and key in weight_dict:
                 loss *= weight_dict[key]
             if "area" in output_dict:
                 loss *= output_dict["area"]

ppsci/loss/l2.py

@@ -87,7 +87,7 @@ def forward(self, output_dict, label_dict, weight_dict=None):
         losses = 0.0
         for key in label_dict:
             loss = F.mse_loss(output_dict[key], label_dict[key], "none")
-            if weight_dict:
+            if weight_dict and key in weight_dict:
                 loss *= weight_dict[key]
 
             if "area" in output_dict:
@@ -181,7 +181,7 @@ def forward(self, output_dict, label_dict, weight_dict=None):
             loss = F.mse_loss(
                 output_dict[key][:n_output], output_dict[key][n_output:], "none"
             )
-            if weight_dict:
+            if weight_dict and key in weight_dict:
                 loss *= weight_dict[key]
 
             if "area" in output_dict:

ppsci/loss/mae.py

@@ -80,7 +80,7 @@ def forward(self, output_dict, label_dict, weight_dict=None):
         losses = 0.0
         for key in label_dict:
             loss = F.l1_loss(output_dict[key], label_dict[key], "none")
-            if weight_dict:
+            if weight_dict and key in weight_dict:
                 loss *= weight_dict[key]
 
             if "area" in output_dict:

ppsci/loss/mse.py

@@ -80,7 +80,7 @@ def forward(self, output_dict, label_dict, weight_dict=None):
         losses = 0.0
         for key in label_dict:
             loss = F.mse_loss(output_dict[key], label_dict[key], "none")
-            if weight_dict:
+            if weight_dict and key in weight_dict:
                 loss *= weight_dict[key]
 
             if "area" in output_dict: