Merge pull request #336 from HydrogenSulfate/add_eval_decorator_2

add eval wrapper to optimize eval code in solver

Author: HydrogenSulfate

ppsci/solver/solver.py

@@ -227,7 +227,9 @@ def __init__(
 
         # decorate model(s) and optimizer(s) for AMP
         if self.use_amp:
-            self.model = amp.decorate(self.model, self.optimizer, self.amp_level)
+            self.model, self.optimizer = amp.decorate(
+                self.model, self.optimizer, self.amp_level
+            )
 
         # wrap model and optimizer to parallel object
         self.rank = dist.get_rank()
@@ -379,6 +381,7 @@ def train(self):
                 )
                 logger.scaler("eval_metric", cur_metric, epoch_id, self.vdl_writer)
 
+                # visualize after evaluation
                 if self.visualizer is not None:
                     self.visualize(epoch_id)
 
@@ -398,7 +401,7 @@ def train(self):
                     self.equation,
                 )
 
-            # always save the latest model for convenient resume training
+            # save the latest model for convenient resume training
             save_load.save_checkpoint(
                 self.model,
                 self.optimizer,
@@ -413,12 +416,9 @@ def train(self):
         if self.vdl_writer is not None:
             self.vdl_writer.close()
 
+    @misc.run_on_eval_mode
     def eval(self, epoch_id: int = 0):
         """Evaluation"""
-        train_state = self.model.training
-        if train_state:
-            self.model.eval()
-
         # set eval func
         self.eval_func = ppsci.solver.eval.eval_func
 
@@ -429,26 +429,19 @@ def eval(self, epoch_id: int = 0):
         logger.info(f"[Eval][Epoch {epoch_id}][Avg] {metric_msg}")
         self.eval_output_info.clear()
 
-        if train_state:
-            self.model.train()
         return result
 
+    @misc.run_on_eval_mode
     def visualize(self, epoch_id: int = 0):
         """Visualization"""
-        train_state = self.model.training
-        if train_state:
-            self.model.eval()
-
         # init train func
         self.visu_func = ppsci.solver.visu.visualize_func
 
         self.visu_func(self, epoch_id)
         logger.info(f"[Visualize][Epoch {epoch_id}] Finished visualization")
 
-        if train_state:
-            self.model.train()
-
     @paddle.no_grad()
+    @misc.run_on_eval_mode
     def predict(
         self,
         input_dict: Dict[str, Union[np.ndarray, paddle.Tensor]],
@@ -463,10 +456,6 @@ def predict(
         Returns:
             Dict[str, paddle.Tensor]: Prediction in dict.
         """
-        train_state = self.model.training
-        if train_state:
-            self.model.eval()
-
         if self.world_size > 1:
             raise NotImplementedError(
                 "Solver.predict only support single device yet, "
@@ -501,10 +490,9 @@ def predict(
 
         pred_dict = {key: paddle.concat(value) for key, value in pred_dict.items()}
 
-        if train_state:
-            self.model.train()
         return pred_dict
 
+    @misc.run_on_eval_mode
     def export(self):
         """Export to inference model"""
         pretrained_path = self.cfg["Global"]["pretrained_model"]

ppsci/utils/misc.py

@@ -13,7 +13,9 @@
 # limitations under the License.
 
 import collections
+import functools
 import random
+from typing import Callable
 from typing import Dict
 from typing import List
 from typing import Tuple
@@ -32,6 +34,7 @@
     "stack_dict_list",
     "combine_array_with_time",
     "set_random_seed",
+    "run_on_eval_mode",
 ]
 
 
@@ -229,3 +232,33 @@ def set_random_seed(seed: int):
     paddle.seed(seed)
     np.random.seed(seed)
     random.seed(seed)
+
+
+def run_on_eval_mode(func: Callable) -> Callable:
+    """A decorator automatically running given class method in eval mode and keep
+    training state unchanged after function finished.
+
+    Args:
+        func (Callable): Class method which is expected running in eval mode.
+
+    Returns:
+        Callable: Decorated class method.
+    """
+
+    @functools.wraps(func)
+    def function_with_eval_state(self, *args, **kwargs):
+        # log original state
+        train_state = self.model.training
+
+        # switch to eval mode
+        if train_state:
+            self.model.eval()
+
+        # run func in eval mode
+        func(self, *args, **kwargs)
+
+        # restore state
+        if train_state:
+            self.model.train()
+
+    return function_with_eval_state