Merge pull request #330 from HydrogenSulfate/opt_code2

optimize code

Author: HydrogenSulfate

README.md

@@ -58,7 +58,7 @@ PaddleScience 是一个基于深度学习框架 PaddlePaddle 开发的科学计
 
 ## 快速开始
 
-参考 [**快速开始**](https://paddlescience-docs.readthedocs.io/zh/latest/zh/quickstart/)
+请参考 [**快速开始**](https://paddlescience-docs.readthedocs.io/zh/latest/zh/quickstart/)
 
 ## 经典案例
 

ppsci/metric/anomaly_coef.py

@@ -63,12 +63,11 @@ def __init__(
         unlog: bool = False,
         scale: float = 1e-5,
     ):
-        super().__init__()
+        super().__init__(keep_batch)
         self.num_lat = num_lat
         self.mean = (
             None if mean is None else paddle.to_tensor(mean, paddle.get_default_dtype())
         )
-        self.keep_batch = keep_batch
         self.variable_dict = variable_dict
         self.unlog = unlog
         self.scale = scale
@@ -110,14 +109,10 @@ def forward(self, output_dict, label_dict):
                     if self.keep_batch:
                         metric_dict[f"{key}.{variable_name}"] = rmse[:, idx]
                     else:
-                        metric_dict[f"{key}.{variable_name}"] = float(
-                            rmse[:, idx].mean()
-                        )
+                        metric_dict[f"{key}.{variable_name}"] = rmse[:, idx].mean()
             else:
                 if self.keep_batch:
-                    rmse = rmse.mean(axis=1)
-                    metric_dict[key] = rmse
+                    metric_dict[key] = rmse.mean(axis=1)
                 else:
-                    rmse = rmse.mean()
-                    metric_dict[key] = float(rmse)
+                    metric_dict[key] = rmse.mean()
         return metric_dict

ppsci/metric/base.py

@@ -18,5 +18,6 @@
 class Metric(nn.Layer):
     """Base class for metric."""
 
-    def __init__(self):
+    def __init__(self, keep_batch: bool = False):
         super().__init__()
+        self.keep_batch = keep_batch

ppsci/metric/l2_rel.py

@@ -24,13 +24,18 @@ class L2Rel(base.Metric):
     metric = \dfrac{\Vert x-y \Vert_2}{\Vert y \Vert_2}
     $$
 
+    Args:
+        keep_batch (bool, optional): Whether keep batch axis. Defaults to False.
+
     Examples:
         >>> import ppsci
         >>> metric = ppsci.metric.L2Rel()
     """
 
-    def __init__(self):
-        super().__init__()
+    def __init__(self, keep_batch: bool = False):
+        if keep_batch:
+            raise ValueError(f"keep_batch should be False, but got {keep_batch}.")
+        super().__init__(keep_batch)
 
     @paddle.no_grad()
     def forward(self, output_dict, label_dict):
@@ -39,6 +44,6 @@ def forward(self, output_dict, label_dict):
             rel_l2 = paddle.norm(label_dict[key] - output_dict[key]) / paddle.norm(
                 label_dict[key]
             )
-            metric_dict[key] = float(rel_l2)
+            metric_dict[key] = rel_l2
 
         return metric_dict

ppsci/metric/mae.py

@@ -34,8 +34,7 @@ class MAE(base.Metric):
     """
 
     def __init__(self, keep_batch: bool = False):
-        super().__init__()
-        self.keep_batch = keep_batch
+        super().__init__(keep_batch)
 
     @paddle.no_grad()
     def forward(self, output_dict, label_dict):
@@ -45,6 +44,6 @@ def forward(self, output_dict, label_dict):
             if self.keep_batch:
                 metric_dict[key] = mae.mean(axis=tuple(range(1, mae.ndim)))
             else:
-                metric_dict[key] = float(mae.mean())
+                metric_dict[key] = mae.mean()
 
         return metric_dict

ppsci/metric/mse.py

@@ -25,19 +25,25 @@ class MSE(base.Metric):
     metric = \dfrac{1}{N}\sum\limits_{i=1}^{N}{(x_i-y_i)^2}
     $$
 
+    Args:
+        keep_batch (bool, optional): Whether keep batch axis. Defaults to False.
+
     Examples:
         >>> import ppsci
         >>> metric = ppsci.metric.MSE()
     """
 
-    def __init__(self):
-        super().__init__()
+    def __init__(self, keep_batch: bool = False):
+        super().__init__(keep_batch)
 
     @paddle.no_grad()
     def forward(self, output_dict, label_dict):
         metric_dict = {}
         for key in label_dict:
-            mse = F.mse_loss(output_dict[key], label_dict[key], "mean")
-            metric_dict[key] = float(mse)
+            mse = F.mse_loss(output_dict[key], label_dict[key], "none")
+            if self.keep_batch:
+                metric_dict[key] = mse.mean(axis=tuple(range(1, mse.ndim)))
+            else:
+                metric_dict[key] = mse.mean()
 
         return metric_dict

ppsci/metric/rmse.py

@@ -31,20 +31,25 @@ class RMSE(base.Metric):
     metric = \sqrt{\dfrac{1}{N}\sum\limits_{i=1}^{N}{(x_i-y_i)^2}}
     $$
 
+    Args:
+        keep_batch (bool, optional): Whether keep batch axis. Defaults to False.
+
     Examples:
         >>> import ppsci
         >>> metric = ppsci.metric.RMSE()
     """
 
-    def __init__(self):
-        super().__init__()
+    def __init__(self, keep_batch: bool = False):
+        if keep_batch:
+            raise ValueError(f"keep_batch should be False, but got {keep_batch}.")
+        super().__init__(keep_batch)
 
     @paddle.no_grad()
     def forward(self, output_dict, label_dict):
         metric_dict = {}
         for key in label_dict:
             rmse = F.mse_loss(output_dict[key], label_dict[key], "mean") ** 0.5
-            metric_dict[key] = float(rmse)
+            metric_dict[key] = rmse
 
         return metric_dict
 
@@ -88,18 +93,16 @@ def __init__(
         unlog: bool = False,
         scale: float = 1e-5,
     ):
-        super().__init__()
+        super().__init__(keep_batch)
         self.num_lat = num_lat
         self.std = (
             None
             if std is None
             else paddle.to_tensor(std, paddle.get_default_dtype()).reshape((1, -1))
         )
-        self.keep_batch = keep_batch
         self.variable_dict = variable_dict
         self.unlog = unlog
         self.scale = scale
-
         self.weight = self.get_latitude_weight(num_lat)
 
     def get_latitude_weight(self, num_lat: int = 720):
@@ -127,18 +130,10 @@ def forward(self, output_dict, label_dict):
                 rmse = rmse * self.std
             if self.variable_dict is not None:
                 for variable_name, idx in self.variable_dict.items():
-                    if self.keep_batch:
-                        metric_dict[f"{key}.{variable_name}"] = rmse[:, idx]
-                    else:
-                        metric_dict[f"{key}.{variable_name}"] = float(
-                            rmse[:, idx].mean()
-                        )
+                    metric_dict[f"{key}.{variable_name}"] = (
+                        rmse[:, idx] if self.keep_batch else rmse[:, idx].mean()
+                    )
             else:
-                if self.keep_batch:
-                    rmse = rmse.mean(axis=1)
-                    metric_dict[key] = rmse
-                else:
-                    rmse = rmse.mean()
-                    metric_dict[key] = float(rmse)
+                metric_dict[key] = rmse.mean(axis=1) if self.keep_batch else rmse.mean()
 
         return metric_dict

ppsci/solver/eval.py

@@ -24,7 +24,7 @@
 
 
 def _eval_by_dataset(solver, epoch_id: int, log_freq: int) -> float:
-    """Evaluation program by dataset.
+    """Evaluate with computing metric on total samples.
 
     Args:
         solver (solver.Solver): Main Solver.
@@ -96,12 +96,12 @@ def _eval_by_dataset(solver, epoch_id: int, log_freq: int) -> float:
             batch_cost = time.perf_counter() - batch_tic
             solver.eval_time_info["reader_cost"].update(reader_cost)
             solver.eval_time_info["batch_cost"].update(batch_cost)
-            total_batch_size = sum([v.shape[0] for v in input_dict.values()])
-            printer.update_eval_loss(solver, loss_dict, total_batch_size)
+            batch_size = next(iter(input_dict.values())).shape[0]
+            printer.update_eval_loss(solver, loss_dict, batch_size)
             if iter_id == 1 or iter_id % log_freq == 0:
                 printer.log_eval_info(
                     solver,
-                    total_batch_size,
+                    batch_size,
                     epoch_id,
                     len(_validator.data_loader),
                     iter_id,
@@ -110,7 +110,7 @@ def _eval_by_dataset(solver, epoch_id: int, log_freq: int) -> float:
             reader_tic = time.perf_counter()
             batch_tic = time.perf_counter()
 
-        # gather all data
+        # concate all data and discard padded sample(s)
         for key in all_input:
             all_input[key] = paddle.concat(all_input[key])
             if len(all_input[key]) > num_samples:
@@ -134,22 +134,22 @@ def _eval_by_dataset(solver, epoch_id: int, log_freq: int) -> float:
                     solver.eval_output_info[metric_str] = misc.AverageMeter(
                         metric_str, ".5f"
                     )
-                solver.eval_output_info[metric_str].update(metric_value, num_samples)
+                solver.eval_output_info[metric_str].update(
+                    float(metric_value), num_samples
+                )
 
+        # use the first metric for return value
         if target_metric is None:
             tmp = metric
             while isinstance(tmp, dict):
                 tmp = next(iter(tmp.values()))
-            assert isinstance(
-                tmp, (int, float)
-            ), f"Target metric({type(tmp)}) should be a number"
-            target_metric = tmp
+            target_metric = float(tmp)
 
     return target_metric
 
 
 def _eval_by_batch(solver, epoch_id: int, log_freq: int) -> float:
-    """Evaluation program by batch.
+    """Evaluate with computing metric by batch, which is memory-efficient.
 
     Args:
         solver (solver.Solver): Main Solver.
@@ -179,7 +179,7 @@ def _eval_by_batch(solver, epoch_id: int, log_freq: int) -> float:
                 for key in solver.eval_time_info:
                     solver.eval_time_info[key].reset()
             reader_cost = time.perf_counter() - reader_tic
-            total_batch_size = next(iter(input_dict.values())).shape[0]
+            batch_size = next(iter(input_dict.values())).shape[0]
 
             for v in input_dict.values():
                 v.stop_gradient = False
@@ -211,11 +211,11 @@ def _eval_by_batch(solver, epoch_id: int, log_freq: int) -> float:
             batch_cost = time.perf_counter() - batch_tic
             solver.eval_time_info["reader_cost"].update(reader_cost)
             solver.eval_time_info["batch_cost"].update(batch_cost)
-            printer.update_eval_loss(solver, loss_dict, total_batch_size)
+            printer.update_eval_loss(solver, loss_dict, batch_size)
             if iter_id == 1 or iter_id % log_freq == 0:
                 printer.log_eval_info(
                     solver,
-                    total_batch_size,
+                    batch_size,
                     epoch_id,
                     len(_validator.data_loader),
                     iter_id,
@@ -224,7 +224,7 @@ def _eval_by_batch(solver, epoch_id: int, log_freq: int) -> float:
             reader_tic = time.perf_counter()
             batch_tic = time.perf_counter()
 
-        # gather all metric
+        # concate all metric and discard metric of padded sample(s)
         for metric_name, metric_dict in metric.items():
             for var_name, metric_value in metric_dict.items():
                 metric_value = paddle.concat(metric_value)[:num_samples]
@@ -237,20 +237,18 @@ def _eval_by_batch(solver, epoch_id: int, log_freq: int) -> float:
                     )
                 solver.eval_output_info[metric_str].update(metric_value, num_samples)
 
+        # use the first metric for return value
         if target_metric is None:
             tmp = metric
             while isinstance(tmp, dict):
                 tmp = next(iter(tmp.values()))
-            assert isinstance(
-                tmp, (int, float)
-            ), f"Target metric({type(tmp)}) should be a number"
             target_metric = tmp
 
     return target_metric
 
 
 def eval_func(solver, epoch_id: int, log_freq: int) -> float:
-    """Evaluation program
+    """Evaluation function.
 
     Args:
         solver (solver.Solver): Main Solver.

ppsci/solver/solver.py

@@ -29,6 +29,7 @@
 import visualdl as vdl
 from packaging import version
 from paddle import amp
+from paddle import jit
 from paddle import nn
 from paddle import optimizer as optim
 from paddle.distributed import fleet
@@ -218,6 +219,9 @@ def __init__(
         # 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
+            if self.update_freq != 1:
+                self.update_freq = 1
+                logger.warning("Set update_freq to to 1 when using L-BFGS optimizer.")
         else:
             self.train_epoch_func = ppsci.solver.train.train_epoch_func
 
@@ -511,11 +515,11 @@ def export(self):
 
         input_spec = copy.deepcopy(self.cfg["Export"]["input_shape"])
         config.replace_shape_with_inputspec_(input_spec)
-        static_model = paddle.jit.to_static(self.model, input_spec=input_spec)
+        static_model = jit.to_static(self.model, input_spec=input_spec)
 
         export_dir = self.cfg["Global"]["save_inference_dir"]
         save_path = os.path.join(export_dir, "inference")
-        paddle.jit.save(static_model, save_path)
+        jit.save(static_model, save_path)
         logger.info(f"The inference model has been exported to {export_dir}")
 
     def autocast_context_manager(self) -> contextlib.AbstractContextManager: