optimize amp context code in train.py/eval.py

Author: HydrogenSulfate

ppsci/solver/eval.py

@@ -67,14 +67,7 @@ def eval_func(solver, epoch_id: int, log_freq: int) -> float:
                 evaluator.add_target_expr(output_formula, output_name)
 
             # forward
-            if solver.use_amp:
-                with amp.auto_cast(level=solver.amp_level):
-                    output_dict = evaluator(input_dict)
-                    validator_loss = _validator.loss(
-                        output_dict, label_dict, weight_dict
-                    )
-                    loss_dict[f"loss({_validator.name})"] = float(validator_loss)
-            else:
+            with solver._autocast_context_manager():
                 output_dict = evaluator(input_dict)
                 validator_loss = _validator.loss(output_dict, label_dict, weight_dict)
                 loss_dict[f"loss({_validator.name})"] = float(validator_loss)

ppsci/solver/solver.py

@@ -14,11 +14,14 @@
 
 from __future__ import annotations
 
+import contextlib
 import copy
 import os
+import sys
 from typing import Any
 from typing import Dict
 from typing import Optional
+from typing import Union
 
 import paddle
 import paddle.amp as amp
@@ -501,3 +504,20 @@ def export(self):
         save_path = os.path.join(export_dir, "inference")
         paddle.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:
+        """Autocast context manager for Auto Mix Precision.
+
+        Returns:
+            Union[contextlib.AbstractContextManager]: Context manager.
+        """
+        if self.use_amp:
+            ctx_manager = amp.auto_cast(level=self.amp_level)
+        else:
+            ctx_manager = (
+                contextlib.nullcontext()
+                if sys.version_info >= (3, 7)
+                else contextlib.suppress()
+            )
+
+        return ctx_manager

ppsci/solver/train.py

@@ -14,9 +14,6 @@
 
 import time
 
-import paddle.amp as amp
-
-from ppsci import solver
 from ppsci.solver import printer
 from ppsci.utils import expression
 from ppsci.utils import misc
@@ -62,14 +59,7 @@ def train_epoch_func(solver, epoch_id: int, log_freq: int):
                 evaluator.add_target_expr(output_formula, output_name)
 
             # forward for every constraint
-            if solver.use_amp:
-                with amp.auto_cast(level=solver.amp_level):
-                    output_dict = evaluator(input_dict)
-                    constraint_loss = _constraint.loss(
-                        output_dict, label_dict, weight_dict
-                    )
-                    total_loss += constraint_loss
-            else:
+            with solver._autocast_context_manager():
                 output_dict = evaluator(input_dict)
                 constraint_loss = _constraint.loss(output_dict, label_dict, weight_dict)
                 total_loss += constraint_loss