Merge branch 'develop' into support_more_fptype

HydrogenSulfate · HydrogenSulfate · commit cceef44059a2 · 2023-05-10T06:37:28.000Z
diff --git a/paddlescience/network/grad_norm.py b/paddlescience/network/grad_norm.py
@@ -84,8 +84,8 @@ def get_grad_norm_loss(self, losses):
         norms = []
         for i in range(losses.shape[0]):
             grad = paddle.autograd.grad(losses[i], W, retain_graph=True)
-            norms.append(paddle.norm(self.loss_weights[i] * grad[0], p=2))
-        norms = paddle.concat(norms)
+            norms.append(paddle.norm(self.loss_weights[i] * grad[0], p=2).reshape([]))
+        norms = paddle.stack(norms)
 
         # calculate the inverse train rate
         loss_ratio = losses.numpy() / self.initial_losses
diff --git a/ppsci/arch/activation.py b/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,
diff --git a/ppsci/arch/base.py b/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
 
diff --git a/ppsci/arch/embedding_koopman.py b/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
 
diff --git a/ppsci/arch/mlp.py b/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
diff --git a/ppsci/data/__init__.py b/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
diff --git a/ppsci/equation/pde/base.py b/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:
diff --git a/ppsci/loss/base.py b/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):
diff --git a/ppsci/metric/base.py b/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):
diff --git a/ppsci/solver/eval.py b/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)
diff --git a/ppsci/solver/solver.py b/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
diff --git a/ppsci/solver/train.py b/ppsci/solver/train.py
@@ -14,8 +14,6 @@
 
 import time
 
-import paddle.amp as amp
-
 from ppsci.solver import printer
 from ppsci.utils import expression
 from ppsci.utils import misc
@@ -61,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
diff --git a/ppsci/solver/visu.py b/ppsci/solver/visu.py
@@ -16,7 +16,6 @@
 import os.path as osp
 
 import paddle
-import paddle.amp as amp
 
 from ppsci.utils import expression
 from ppsci.utils import misc
@@ -64,10 +63,7 @@ def visualize_func(solver, epoch_id: int):
                 evaluator.add_target_expr(output_expr, output_key)
 
             # forward
-            if solver.use_amp:
-                with amp.auto_cast(level=solver.amp_level):
-                    batch_output_dict = evaluator(batch_input_dict)
-            else:
+            with solver._autocast_context_manager():
                 batch_output_dict = evaluator(batch_input_dict)
 
             # collect batch data
diff --git a/ppsci/utils/expression.py b/ppsci/utils/expression.py
@@ -154,8 +154,8 @@ def add_target_expr(self, expr: Callable, expr_name: str):
         """Add an expression `expr` named `expr_name` to
 
         Args:
-            expr (Callable): _description_
-            expr_name (str): _description_
+            expr (Callable): Callable function for computing an expression.
+            expr_name (str): Name of expression.
         """
         self.expr_dict[expr_name] = expr
 
diff --git a/ppsci/utils/initializer.py b/ppsci/utils/initializer.py
@@ -25,7 +25,7 @@
 
 import numpy as np
 import paddle
-import paddle.nn as nn
+from paddle import nn
 from typing_extensions import Literal
 
 from ppsci.utils import logger
diff --git a/ppsci/utils/logger.py b/ppsci/utils/logger.py
@@ -158,15 +158,15 @@ def advertise():
     ===========================================================
 
     """
-    copyright = "PaddleScience is powered by PaddlePaddle !"
+    _copyright = "PaddleScience is powered by PaddlePaddle !"
     ad = "For more info please go to the following website."
     website = "https://github.com/PaddlePaddle/PaddleScience"
-    AD_LEN = 6 + len(max([copyright, ad, website], key=len))
+    AD_LEN = 6 + len(max([_copyright, ad, website], key=len))
 
     info(
         "\n{0}\n{1}\n{2}\n{3}\n{4}\n{5}\n{6}\n{7}\n".format(
             "=" * (AD_LEN + 4),
-            "=={}==".format(copyright.center(AD_LEN)),
+            "=={}==".format(_copyright.center(AD_LEN)),
             "=" * (AD_LEN + 4),
             "=={}==".format(" " * AD_LEN),
             "=={}==".format(ad.center(AD_LEN)),
diff --git a/ppsci/utils/reader.py b/ppsci/utils/reader.py
@@ -14,6 +14,7 @@
 
 import collections
 import csv
+import sys
 from typing import Dict
 from typing import Optional
 from typing import Tuple
@@ -56,7 +57,7 @@ def load_csv_file(
                     raw_data[key].append(value)
     except FileNotFoundError:
         logger.error(f"{file_path} isn't a valid csv file.")
-        exit(0)
+        sys.exit()
 
     # convert to numpy array
     data_dict = {}
