add OptimizerList for multiple optimizer case (#348)

* add OptimizerList for multiple optimizer case

* refine docstring and type-hint

* refine solver

* remoe export code for it is not supported yet

* remove redundant f"

* refine docstring and add OptimizerList to api doc

* refine code with flake8

Author: HydrogenSulfate

docs/zh/api/optimizer.md

@@ -10,5 +10,6 @@
         - AdamW
         - RMSProp
         - LBFGS
+        - OptimizerList
       show_root_heading: false
       heading_level: 3

ppsci/geometry/mesh.py

@@ -48,7 +48,7 @@ def __init__(self, mesh: Union["pymesh.Mesh", str]):
         elif isinstance(mesh, pymesh.Mesh):
             self.py_mesh = mesh
         else:
-            raise ValueError(f"arg `mesh` should be path string or or `pymesh.Mesh`")
+            raise ValueError("arg `mesh` should be path string or or `pymesh.Mesh`")
 
         self.init_mesh()
 

ppsci/optimizer/__init__.py

@@ -20,9 +20,19 @@
 from ppsci.optimizer.optimizer import Adam
 from ppsci.optimizer.optimizer import AdamW
 from ppsci.optimizer.optimizer import Momentum
+from ppsci.optimizer.optimizer import OptimizerList
 from ppsci.optimizer.optimizer import RMSProp
 
-__all__ = ["LBFGS", "SGD", "Adam", "AdamW", "Momentum", "RMSProp", "lr_scheduler"]
+__all__ = [
+    "LBFGS",
+    "SGD",
+    "Adam",
+    "AdamW",
+    "Momentum",
+    "RMSProp",
+    "OptimizerList",
+    "lr_scheduler",
+]
 
 
 def build_lr_scheduler(cfg, epochs, iters_per_epoch):

ppsci/optimizer/optimizer.py

@@ -12,19 +12,26 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import TYPE_CHECKING
+from typing import Dict
+from typing import List
 from typing import Optional
 from typing import Tuple
 from typing import Union
 
+if TYPE_CHECKING:
+    import paddle
+
 from paddle import nn
 from paddle import optimizer as optim
 from paddle import regularizer
 from paddle.incubate import optimizer as incubate_optim
 from typing_extensions import Literal
 
 from ppsci.utils import logger
+from ppsci.utils import misc
 
-__all__ = ["SGD", "Momentum", "Adam", "RMSProp", "AdamW", "LBFGS"]
+__all__ = ["SGD", "Momentum", "Adam", "RMSProp", "AdamW", "LBFGS", "OptimizerList"]
 
 
 class SGD:
@@ -452,3 +459,55 @@ def __call__(self, model_list: Tuple[nn.Layer, ...]):
 
     def _apply_decay_param_fun(self, name):
         return name not in self.no_weight_decay_param_name_list
+
+
+class OptimizerList:
+    """OptimizerList which wrap more than one optimizer.
+    NOTE: LBFGS is not supported yet.
+
+    Args:
+        optimizer_list (Tuple[optim.Optimizer, ...]): Optimizers listed in a tuple.
+
+    Examples:
+        >>> import ppsci
+        >>> model1 = ppsci.arch.MLP(("x",), ("u",), 5, 20)
+        >>> opt1 = ppsci.optimizer.Adam(1e-3)((model1,))
+        >>> model2 = ppsci.arch.MLP(("y",), ("v",), 5, 20)
+        >>> opt2 = ppsci.optimizer.Adam(1e-3)((model2,))
+        >>> opt = ppsci.optimizer.OptimizerList((opt1, opt2))
+    """
+
+    def __init__(self, optimizer_list: Tuple[optim.Optimizer, ...]):
+        super().__init__()
+        self._opt_list = optimizer_list
+        if "LBFGS" in set(misc.typename(opt) for opt in optimizer_list):
+            raise ValueError("LBFGS is not supported in OptimizerList yet.")
+
+    def step(self):
+        for opt in self._opt_list:
+            opt.step()
+
+    def clear_grad(self):
+        for opt in self._opt_list:
+            opt.clear_grad()
+
+    def get_lr(self) -> float:
+        """Return learning rate of first optimizer"""
+        return self._opt_list[0].get_lr()
+
+    def set_state_dict(self, state_dicts: List[Dict[str, "paddle.Tensor"]]):
+        for i, opt in enumerate(self._opt_list):
+            opt.set_state_dict(state_dicts[i])
+
+    def state_dict(self) -> List[Dict[str, "paddle.Tensor"]]:
+        state_dicts = [opt.state_dict() for opt in self._opt_list]
+        return state_dicts
+
+    def __len__(self) -> int:
+        return len(self._opt_list)
+
+    def __getitem__(self, idx):
+        return self._opt_list[idx]
+
+    def __setitem__(self, idx, opt):
+        raise NotImplementedError("Can not modify any item in OptimizerList.")

ppsci/solver/eval.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import time
+from typing import TYPE_CHECKING
 
 import paddle
 from paddle import io
@@ -21,8 +22,11 @@
 from ppsci.utils import misc
 from ppsci.utils import profiler
 
+if TYPE_CHECKING:
+    from ppsci import solver
 
-def _eval_by_dataset(solver, epoch_id: int, log_freq: int) -> float:
+
+def _eval_by_dataset(solver: "solver.Solver", epoch_id: int, log_freq: int) -> float:
     """Evaluate with computing metric on total samples.
 
     Args:
@@ -148,7 +152,7 @@ def _eval_by_dataset(solver, epoch_id: int, log_freq: int) -> float:
     return target_metric
 
 
-def _eval_by_batch(solver, epoch_id: int, log_freq: int) -> float:
+def _eval_by_batch(solver: "solver.Solver", epoch_id: int, log_freq: int) -> float:
     """Evaluate with computing metric by batch, which is memory-efficient.
 
     Args:
@@ -249,7 +253,7 @@ def _eval_by_batch(solver, epoch_id: int, log_freq: int) -> float:
     return target_metric
 
 
-def eval_func(solver, epoch_id: int, log_freq: int) -> float:
+def eval_func(solver: "solver.Solver", epoch_id: int, log_freq: int) -> float:
     """Evaluation function.
 
     Args:

ppsci/solver/solver.py

@@ -15,7 +15,6 @@
 from __future__ import annotations
 
 import contextlib
-import copy
 import itertools
 import os
 import sys
@@ -51,7 +50,7 @@ class Solver:
     Args:
         model (nn.Layer): Model.
         constraint (Optional[Dict[str, ppsci.constraint.Constraint]]): Constraint(s) applied on model. Defaults to None.
-        output_dir (str, optional): Output directory. Defaults to "./output/".
+        output_dir (Optional[str]): Output directory. Defaults to "./output/".
         optimizer (Optional[optimizer.Optimizer]): Optimizer object. Defaults to None.
         lr_scheduler (Optional[optimizer.lr.LRScheduler]): Learning rate scheduler. Defaults to None.
         epochs (int, optional): Training epoch(s). Defaults to 5.
@@ -108,7 +107,7 @@ def __init__(
         self,
         model: nn.Layer,
         constraint: Optional[Dict[str, ppsci.constraint.Constraint]] = None,
-        output_dir: str = "./output/",
+        output_dir: Optional[str] = "./output/",
         optimizer: Optional[optim.Optimizer] = None,
         lr_scheduler: Optional[optim.lr.LRScheduler] = None,
         epochs: int = 5,
@@ -371,7 +370,7 @@ def from_config(cfg: Dict[str, Any]) -> Solver:
         )
 
     def train(self):
-        """Training"""
+        """Training."""
         self.global_step = self.best_metric["epoch"] * self.iters_per_epoch + 1
 
         for epoch_id in range(self.best_metric["epoch"] + 1, self.epochs + 1):
@@ -446,8 +445,15 @@ def train(self):
             self.vdl_writer.close()
 
     @misc.run_on_eval_mode
-    def eval(self, epoch_id: int = 0):
-        """Evaluation"""
+    def eval(self, epoch_id: int = 0) -> float:
+        """Evaluation.
+
+        Args:
+            epoch_id (int, optional): Epoch id. Defaults to 0.
+
+        Returns:
+            float: The value of the evaluation, used to judge the quality of the model.
+        """
         # set eval func
         self.eval_func = ppsci.solver.eval.eval_func
 
@@ -462,8 +468,12 @@ def eval(self, epoch_id: int = 0):
 
     @misc.run_on_eval_mode
     def visualize(self, epoch_id: int = 0):
-        """Visualization"""
-        # init train func
+        """Visualization.
+
+        Args:
+            epoch_id (int, optional): Epoch id. Defaults to 0.
+        """
+        # set visualize func
         self.visu_func = ppsci.solver.visu.visualize_func
 
         self.visu_func(self, epoch_id)
@@ -568,21 +578,8 @@ def predict(
 
     @misc.run_on_eval_mode
     def export(self):
-        """Export to inference model"""
-        pretrained_path = self.cfg["Global"]["pretrained_model"]
-        if pretrained_path is not None:
-            save_load.load_pretrain(self.model, pretrained_path, self.equation)
-
-        self.model.eval()
-
-        input_spec = copy.deepcopy(self.cfg["Export"]["input_shape"])
-        config.replace_shape_with_inputspec_(input_spec)
-        static_model = jit.to_static(self.model, input_spec=input_spec)
-
-        export_dir = self.cfg["Global"]["save_inference_dir"]
-        save_path = os.path.join(export_dir, "inference")
-        jit.save(static_model, save_path)
-        logger.info(f"The inference model has been exported to {export_dir}")
+        """Export to inference model."""
+        raise NotImplementedError("model export is not supported yet.")
 
     def autocast_context_manager(
         self, enable: bool, level: Literal["O0", "O1", "O2"] = "O1"

ppsci/solver/train.py

@@ -13,10 +13,13 @@
 # limitations under the License.
 
 import time
+from typing import TYPE_CHECKING
 
 from paddle.distributed.fleet.utils import hybrid_parallel_util as hpu
 
-from ppsci import solver
+if TYPE_CHECKING:
+    from ppsci import solver
+
 from ppsci.solver import printer
 from ppsci.utils import misc
 from ppsci.utils import profiler

ppsci/utils/expression.py

@@ -18,19 +18,20 @@
 from typing import Optional
 from typing import Tuple
 
-import paddle
 from paddle import jit
 from paddle import nn
 
 if TYPE_CHECKING:
+    import paddle
     from ppsci import constraint
     from ppsci import validate
 
 from ppsci.autodiff import clear
 
 
 class ExpressionSolver(nn.Layer):
-    """Expression Solver
+    """Expression computing helper, which compute named result according to corresponding
+    function and related inputs.
 
     Examples:
         >>> import ppsci
@@ -45,12 +46,26 @@ def __init__(self):
     def train_forward(
         self,
         expr_dicts: Tuple[Dict[str, Callable], ...],
-        input_dicts: Tuple[Dict[str, paddle.Tensor], ...],
+        input_dicts: Tuple[Dict[str, "paddle.Tensor"], ...],
         model: nn.Layer,
         constraint: Dict[str, "constraint.Constraint"],
-        label_dicts: Tuple[Dict[str, paddle.Tensor], ...],
-        weight_dicts: Tuple[Dict[str, paddle.Tensor], ...],
-    ):
+        label_dicts: Tuple[Dict[str, "paddle.Tensor"], ...],
+        weight_dicts: Tuple[Dict[str, "paddle.Tensor"], ...],
+    ) -> Tuple["paddle.Tensor", ...]:
+        """Forward computation for training, including model forward and equation
+        forward.
+
+        Args:
+            expr_dicts (Tuple[Dict[str, Callable], ...]): Tuple of expression dicts.
+            input_dicts (Tuple[Dict[str, paddle.Tensor], ...]): Tuple of input dicts.
+            model (nn.Layer): NN model.
+            constraint (Dict[str, "constraint.Constraint"]): Constraint dict.
+            label_dicts (Tuple[Dict[str, paddle.Tensor], ...]): Tuple of label dicts.
+            weight_dicts (Tuple[Dict[str, paddle.Tensor], ...]): Tuple of weight dicts.
+
+        Returns:
+            Tuple[paddle.Tensor, ...]: Tuple of losses for each constraint.
+        """
         output_dicts = []
         for i, expr_dict in enumerate(expr_dicts):
             # model forward
@@ -90,12 +105,27 @@ def train_forward(
     def eval_forward(
         self,
         expr_dict: Dict[str, Callable],
-        input_dict: Dict[str, paddle.Tensor],
+        input_dict: Dict[str, "paddle.Tensor"],
         model: nn.Layer,
         validator: "validate.Validator",
-        label_dict: Dict[str, paddle.Tensor],
-        weight_dict: Dict[str, paddle.Tensor],
-    ):
+        label_dict: Dict[str, "paddle.Tensor"],
+        weight_dict: Dict[str, "paddle.Tensor"],
+    ) -> Tuple[Dict[str, "paddle.Tensor"], "paddle.Tensor"]:
+        """Forward computation for evaluation, including model forward and equation
+        forward.
+
+        Args:
+            expr_dict (Dict[str, Callable]): Expression dict.
+            input_dict (Dict[str, paddle.Tensor]): Input dict.
+            model (nn.Layer): NN model.
+            validator (validate.Validator): Validator.
+            label_dict (Dict[str, paddle.Tensor]): Label dict.
+            weight_dict (Dict[str, paddle.Tensor]): Weight dict.
+
+        Returns:
+            Tuple[Dict[str, paddle.Tensor], paddle.Tensor]: Result dict and loss for
+                given validator.
+        """
         # model forward
         if callable(next(iter(expr_dict.values()))):
             output_dict = model(input_dict)
@@ -123,9 +153,20 @@ def eval_forward(
     def visu_forward(
         self,
         expr_dict: Optional[Dict[str, Callable]],
-        input_dict: Dict[str, paddle.Tensor],
+        input_dict: Dict[str, "paddle.Tensor"],
         model: nn.Layer,
-    ):
+    ) -> Dict[str, "paddle.Tensor"]:
+        """Forward computation for visualization, including model forward and equation
+        forward.
+
+        Args:
+            expr_dict (Optional[Dict[str, Callable]]): Expression dict.
+            input_dict (Dict[str, paddle.Tensor]]): Input dict.
+            model (nn.Layer): NN model.
+
+        Returns:
+            Dict[str, paddle.Tensor]: Result dict for given expression dict.
+        """
         # model forward
         output_dict = model(input_dict)
 
@@ -140,5 +181,4 @@ def visu_forward(
             # clear differentiation cache
             clear()
 
-        # compute loss for each validator according to its' own output, label and weight
         return output_dict

ppsci/utils/reader.py

@@ -18,7 +18,6 @@
 from typing import Dict
 from typing import Optional
 from typing import Tuple
-from typing import Union
 
 import meshio
 import numpy as np