fix typehint for Solver.predict

Author: HydrogenSulfate

ppsci/solver/solver.py

@@ -21,7 +21,9 @@
 from typing import Any
 from typing import Dict
 from typing import Optional
+from typing import Union
 
+import numpy as np
 import paddle
 import paddle.distributed as dist
 import visualdl as vdl
@@ -67,6 +69,9 @@ class Solver:
         amp_level (Literal["O1", "O2", "O0"], optional): AMP level. Defaults to "O0".
         pretrained_model_path (Optional[str]): Pretrained model path. Defaults to None.
         checkpoint_path (Optional[str]): Checkpoint path. Defaults to None.
+        compute_metric_by_batch (bool, optional): Whether calculate metrics after each batch during evaluate. Defaults to False.
+        eval_with_no_grad (bool, optional): Whether set `stop_gradient=True` for every Tensor if no differentiation
+            involved during computation, generally for save GPU memory and accelerate computing. Defaults to False.
 
     Examples:
         >>> import ppsci
@@ -120,6 +125,8 @@ def __init__(
         amp_level: Literal["O1", "O2", "O0"] = "O0",
         pretrained_model_path: Optional[str] = None,
         checkpoint_path: Optional[str] = None,
+        compute_metric_by_batch: bool = False,
+        eval_with_no_grad: bool = False,
     ):
         # set model
         self.model = model
@@ -190,6 +197,11 @@ def __init__(
         if pretrained_model_path is not None:
             save_load.load_pretrain(self.model, pretrained_model_path, self.equation)
 
+        # whether calculate metrics after each batch during evaluate
+        self.compute_metric_by_batch = compute_metric_by_batch
+        # whether set `stop_gradient=True` for every Tensor if no differentiation involved during computation
+        self.eval_with_no_grad = eval_with_no_grad
+
         # initialize an dict for tracking best metric during training
         self.best_metric = {
             "metric": float("inf"),
@@ -291,6 +303,8 @@ def from_config(cfg: Dict[str, Any]) -> Solver:
         update_freq = cfg["Global"].get("update_freq", 1)
         pretrained_model_path = cfg["Global"].get("pretrained_model_path", None)
         checkpoint_path = cfg["Global"].get("checkpoint_path", None)
+        compute_metric_by_batch = cfg["Global"].get("compute_metric_by_batch", False)
+        eval_with_no_grad = cfg["Global"].get("eval_with_no_grad", False)
 
         return Solver(
             model,
@@ -317,6 +331,8 @@ def from_config(cfg: Dict[str, Any]) -> Solver:
             amp_level,
             pretrained_model_path,
             checkpoint_path,
+            compute_metric_by_batch,
+            eval_with_no_grad,
         )
 
     def train(self):
@@ -430,12 +446,14 @@ def visualize(self, epoch_id: int = 0):
 
     @paddle.no_grad()
     def predict(
-        self, input_dict: Dict[str, paddle.Tensor], batch_size: int = 64
+        self,
+        input_dict: Dict[str, Union[np.ndarray, paddle.Tensor]],
+        batch_size: int = 64,
     ) -> Dict[str, paddle.Tensor]:
         """Pure prediction using model.forward(...), support single device prediction yet.
 
         Args:
-            input_dict (Dict[str, paddle.Tensor]): Input data in dict.
+            input_dict (Dict[str, Union[np.ndarray, paddle.Tensor]]): Input data in dict.
             batch_size (int, optional): Predicting by batch size. Defaults to 64.
 
         Returns:
@@ -516,3 +534,20 @@ def autocast_context_manager(self) -> contextlib.AbstractContextManager:
             )
 
         return ctx_manager
+
+    def no_grad_context_manager(self) -> contextlib.AbstractContextManager:
+        """No grad manager.
+
+        Returns:
+            Union[contextlib.AbstractContextManager]: Context manager.
+        """
+        if self.eval_with_no_grad:
+            ctx_manager = paddle.no_grad()
+        else:
+            ctx_manager = (
+                contextlib.nullcontext()
+                if sys.version_info >= (3, 7)
+                else contextlib.suppress()
+            )
+
+        return ctx_manager