Merge pull request #290 from HydrogenSulfate/opt_train_amp

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

Author: HydrogenSulfate

ppsci/arch/activation.py

@@ -15,8 +15,8 @@
 from typing import Callable
 
 import paddle
-import paddle.nn as nn
 import paddle.nn.functional as F
+from paddle import nn
 
 act_func_dict = {
     "elu": F.elu,

ppsci/arch/base.py

@@ -17,7 +17,7 @@
 
 import numpy as np
 import paddle
-import paddle.nn as nn
+from paddle import nn
 
 from ppsci.utils import logger
 

ppsci/arch/embedding_koopman.py

@@ -21,7 +21,7 @@
 
 import numpy as np
 import paddle
-import paddle.nn as nn
+from paddle import nn
 from paddle.nn.initializer import Constant
 from paddle.nn.initializer import Uniform
 

ppsci/arch/mlp.py

@@ -15,7 +15,7 @@
 from typing import Tuple
 from typing import Union
 
-import paddle.nn as nn
+from paddle import nn
 
 from ppsci.arch import activation as act_mod
 from ppsci.arch import base

ppsci/data/__init__.py

@@ -17,9 +17,9 @@
 from functools import partial
 
 import numpy as np
-import paddle.device as device
 import paddle.distributed as dist
-import paddle.io as io
+from paddle import device
+from paddle import io
 
 from ppsci.data import dataloader
 from ppsci.data import dataset

ppsci/equation/pde/base.py

@@ -18,8 +18,8 @@
 from typing import Tuple
 
 import paddle
-import paddle.nn as nn
 import sympy
+from paddle import nn
 
 
 class PDE:

ppsci/loss/base.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import paddle.nn as nn
+from paddle import nn
 
 
 class Loss(nn.Layer):

ppsci/metric/base.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import paddle.nn as nn
+from paddle import nn
 
 
 class Metric(nn.Layer):

ppsci/solver/eval.py

@@ -15,8 +15,7 @@
 import time
 
 import paddle
-import paddle.amp as amp
-import paddle.io as io
+from paddle import io
 
 from ppsci.solver import printer
 from ppsci.utils import expression
@@ -67,14 +66,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,20 +14,22 @@
 
 from __future__ import annotations
 
+import contextlib
 import copy
 import os
+import sys
 from typing import Any
 from typing import Dict
 from typing import Optional
 
 import paddle
-import paddle.amp as amp
 import paddle.distributed as dist
-import paddle.incubate as incubate
-import paddle.nn as nn
-import paddle.optimizer as optimizer
 import visualdl as vdl
 from packaging import version
+from paddle import amp
+from paddle import incubate
+from paddle import nn
+from paddle import optimizer
 from paddle.distributed import fleet
 from typing_extensions import Literal
 
@@ -57,7 +59,7 @@ class Solver:
         eval_freq (int, optional): Evaluation frequency. Defaults to 1.
         seed (int, optional): Random seed. Defaults to 42.
         vdl_writer (Optional[vdl.LogWriter]): VisualDL writer object. Defaults to None.
-        device (Literal["cpu", "gpu", "xpu"], optional): _description_. Defaults to "gpu".
+        device (Literal["cpu", "gpu", "xpu"], optional): Runtime device. Defaults to "gpu".
         equation (Optional[Dict[str, ppsci.equation.PDE]]): Equation dict. Defaults to None.
         geom (Optional[Dict[str, ppsci.geometry.Geometry]]): Geometry dict. Defaults to None.
         validator (Optional[Dict[str, ppsci.validate.Validator]]): Validator dict. Defaults to None.
@@ -469,10 +471,7 @@ def predict(
                 batch_input_dict[key].stop_gradient = False
 
             # forward
-            if self.use_amp:
-                with amp.auto_cast(level=self.amp_level):
-                    batch_output_dict = self.model(batch_input_dict)
-            else:
+            with self._autocast_context_manager():
                 batch_output_dict = self.model(batch_input_dict)
 
             # collect batch data
@@ -501,3 +500,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