refine type hint of train/eval/visu

Author: HydrogenSulfate

ppsci/solver/eval.py

@@ -18,13 +18,14 @@
 import paddle.amp as amp
 import paddle.io as io
 
+from ppsci import solver
 from ppsci.solver import printer
 from ppsci.utils import expression
 from ppsci.utils import misc
 from ppsci.utils import profiler
 
 
-def eval_func(solver, epoch_id, log_freq) -> float:
+def eval_func(solver: solver.Solver, epoch_id: int, log_freq: int) -> float:
     """Evaluation program
 
     Args:

ppsci/solver/solver.py

@@ -392,7 +392,7 @@ def train(self):
         if self.vdl_writer is not None:
             self.vdl_writer.close()
 
-    def eval(self, epoch_id=0):
+    def eval(self, epoch_id: int = 0):
         """Evaluation"""
         train_state = self.model.training
         if train_state:
@@ -412,7 +412,7 @@ def eval(self, epoch_id=0):
             self.model.train()
         return result
 
-    def visualize(self, epoch_id=0):
+    def visualize(self, epoch_id: int = 0):
         """Visualization"""
         train_state = self.model.training
         if train_state:
@@ -461,7 +461,9 @@ def predict(
             # prepare batch input dict
             for key in input_dict:
                 if not paddle.is_tensor(input_dict[key]):
-                    batch_input_dict[key] = paddle.to_tensor(input_dict[key][st:ed])
+                    batch_input_dict[key] = paddle.to_tensor(
+                        input_dict[key][st:ed], paddle.get_default_dtype()
+                    )
                 else:
                     batch_input_dict[key] = input_dict[key][st:ed]
                 batch_input_dict[key].stop_gradient = False

ppsci/solver/train.py

@@ -16,17 +16,18 @@
 
 import paddle.amp as amp
 
+from ppsci import solver
 from ppsci.solver import printer
 from ppsci.utils import expression
 from ppsci.utils import misc
 from ppsci.utils import profiler
 
 
-def train_epoch_func(solver, epoch_id, log_freq):
+def train_epoch_func(solver: solver.Solver, epoch_id: int, log_freq: int):
     """Train program for one epoch
 
     Args:
-        solver (Solver): Main solver.
+        solver (solver.Solver): Main solver.
         epoch_id (int): Epoch id.
         log_freq (int): Log training information every `log_freq` steps.
     """
@@ -112,11 +113,11 @@ def train_epoch_func(solver, epoch_id, log_freq):
         batch_tic = time.perf_counter()
 
 
-def train_LBFGS_epoch_func(solver, epoch_id, log_freq):
+def train_LBFGS_epoch_func(solver: solver.Solver, epoch_id: int, log_freq: int):
     """Train function for one epoch with L-BFGS optimizer.
 
     Args:
-        solver (Solver): Main solver.
+        solver (solver.Solver): Main solver.
         epoch_id (int): Epoch id.
         log_freq (int): Log training information every `log_freq` steps.
     """

ppsci/solver/visu.py

@@ -18,15 +18,17 @@
 import paddle
 import paddle.amp as amp
 
+from ppsci import solver
 from ppsci.utils import expression
 from ppsci.utils import misc
 
 
-def visualize_func(solver, epoch_id):
+@paddle.no_grad()
+def visualize_func(solver: solver.Solver, epoch_id: int):
     """Visualization program
 
     Args:
-        solver (Solver): Main Solver.
+        solver (solver.Solver): Main Solver.
         epoch_id (int): Epoch id.
 
     Returns:
@@ -49,7 +51,9 @@ def visualize_func(solver, epoch_id):
             # prepare batch input dict
             for key in input_dict:
                 if not paddle.is_tensor(input_dict[key]):
-                    batch_input_dict[key] = paddle.to_tensor(input_dict[key][st:ed])
+                    batch_input_dict[key] = paddle.to_tensor(
+                        input_dict[key][st:ed], paddle.get_default_dtype()
+                    )
                 else:
                     batch_input_dict[key] = input_dict[key][st:ed]
                 batch_input_dict[key].stop_gradient = False