[Feature] Add loss aggregator to saving/loading process (#949)

* fix document and formulations

* correct doc of Data.batch_transform

* update schema for None batch size in INFER config

* support batch_size to be set to None where certain data is unnecessary to be batchified

* update pre-commit

* refine EVAL config schema

* save&load state_dict for loss aggregator

Author: HydrogenSulfate

docs/zh/install_setup.md

@@ -95,7 +95,7 @@
 
         ``` sh
         cd PaddleScience/
-        set PYTHONPATH=%cd%
+        set PYTHONPATH=%PYTHONPATH%;%CD%
         pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple # manually install requirements
         ```
 

ppsci/loss/mtl/agda.py

@@ -14,6 +14,7 @@
 
 from __future__ import annotations
 
+from typing import ClassVar
 from typing import List
 
 import paddle
@@ -30,6 +31,10 @@ class AGDA(base.LossAggregator):
 
     NOTE: This loss aggregator is only suitable for two-task learning and the first task loss must be PDE loss.
 
+    Attributes:
+        should_persist(bool): Whether to persist the loss aggregator when saving.
+            Those loss aggregators with parameters and/or buffers should be persisted.
+
     Args:
         model (nn.Layer): Training model.
         M (int, optional): Smoothing period. Defaults to 100.
@@ -49,6 +54,7 @@ class AGDA(base.LossAggregator):
         ...     bc_loss = paddle.sum((y2 - 2) ** 2)
         ...     loss_aggregator({'pde_loss': pde_loss, 'bc_loss': bc_loss}).backward()
     """
+    should_persist: ClassVar[bool] = False
 
     def __init__(self, model: nn.Layer, M: int = 100, gamma: float = 0.999) -> None:
         super().__init__(model)

ppsci/loss/mtl/base.py

@@ -15,6 +15,7 @@
 from __future__ import annotations
 
 from typing import TYPE_CHECKING
+from typing import ClassVar
 from typing import Dict
 from typing import Union
 
@@ -27,10 +28,16 @@
 class LossAggregator(nn.Layer):
     """Base class of loss aggregator mainly for multitask learning.
 
+    Attributes:
+        should_persist(bool): Whether to persist the loss aggregator when saving.
+            Those loss aggregators with parameters and/or buffers should be persisted.
+
     Args:
         model (nn.Layer): Training model.
     """
 
+    should_persist: ClassVar[bool] = False
+
     def __init__(self, model: nn.Layer) -> None:
         super().__init__()
         self.model = model
@@ -52,3 +59,10 @@ def backward(self) -> None:
         raise NotImplementedError(
             f"'backward' should be implemented in subclass {self.__class__.__name__}"
         )
+
+    def state_dict(self):
+        agg_state = super().state_dict()
+        model_state = self.model.state_dictq()
+        # remove model parameters from state dict for already in pdparams
+        agg_state = {k: v for k, v in agg_state.items() if k not in model_state}
+        return agg_state

ppsci/loss/mtl/grad_norm.py

@@ -14,6 +14,7 @@
 
 from __future__ import annotations
 
+from typing import ClassVar
 from typing import Dict
 from typing import List
 
@@ -42,6 +43,10 @@ class GradNorm(base.LossAggregator):
     \end{align*}
     $$
 
+    Attributes:
+        should_persist(bool): Whether to persist the loss aggregator when saving.
+            Those loss aggregators with parameters and/or buffers should be persisted.
+
     Args:
         model (nn.Layer): Training model.
         num_losses (int, optional): Number of losses. Defaults to 1.
@@ -63,6 +68,7 @@ class GradNorm(base.LossAggregator):
         ...     loss2 = paddle.sum((y2 - 2) ** 2)
         ...     loss_aggregator({'loss1': loss1, 'loss2': loss2}).backward()
     """
+    should_persist: ClassVar[bool] = True
     weight: paddle.Tensor
 
     def __init__(

ppsci/loss/mtl/ntk.py

@@ -14,15 +14,21 @@
 
 from __future__ import annotations
 
+from typing import TYPE_CHECKING
+from typing import ClassVar
 from typing import List
 
 import paddle
-from paddle import nn
 
 from ppsci.loss.mtl import base
 
+if TYPE_CHECKING:
+    from paddle import nn
+
 
 class NTK(base.LossAggregator):
+    should_persist: ClassVar[bool] = True
+
     def __init__(
         self,
         model: nn.Layer,

ppsci/loss/mtl/pcgrad.py

@@ -14,6 +14,7 @@
 
 from __future__ import annotations
 
+from typing import ClassVar
 from typing import List
 
 import numpy as np
@@ -31,6 +32,10 @@ class PCGrad(base.LossAggregator):
 
     Code reference: [https://github.com/tianheyu927/PCGrad/blob/master/PCGrad_tf.py](https://github.com/tianheyu927/PCGrad/blob/master/PCGrad_tf.py)
 
+    Attributes:
+        should_persist(bool): Whether to persist the loss aggregator when saving.
+            Those loss aggregators with parameters and/or buffers should be persisted.
+
     Args:
         model (nn.Layer): Training model.
 
@@ -48,6 +53,7 @@ class PCGrad(base.LossAggregator):
         ...     loss2 = paddle.sum((y2 - 2) ** 2)
         ...     loss_aggregator({'loss1': loss1, 'loss2': loss2}).backward()
     """
+    should_persist: ClassVar[bool] = False
 
     def __init__(self, model: nn.Layer) -> None:
         super().__init__(model)

ppsci/loss/mtl/relobralo.py

@@ -14,6 +14,7 @@
 
 from __future__ import annotations
 
+from typing import ClassVar
 from typing import Dict
 
 import paddle
@@ -26,6 +27,10 @@ class Relobralo(nn.Layer):
 
     [Multi-Objective Loss Balancing for Physics-Informed Deep Learning](https://arxiv.org/abs/2110.09813)
 
+    Attributes:
+        should_persist(bool): Whether to persist the loss aggregator when saving.
+            Those loss aggregators with parameters and/or buffers should be persisted.
+
     Args:
         num_losses (int): Number of losses.
         alpha (float, optional): Ability for remembering past in paper. Defaults to 0.95.
@@ -49,6 +54,7 @@ class Relobralo(nn.Layer):
         ...     loss2 = paddle.sum((y2 - 2) ** 2)
         ...     loss_aggregator({'loss1': loss1, 'loss2': loss2}).backward()
     """
+    should_persist: ClassVar[bool] = True
 
     def __init__(
         self,

ppsci/loss/mtl/sum.py

@@ -20,6 +20,8 @@
 if TYPE_CHECKING:
     import paddle
 
+from typing import ClassVar
+
 from ppsci.loss.mtl.base import LossAggregator
 
 
@@ -30,7 +32,12 @@ class Sum(LossAggregator):
     $$
     loss = \sum_i^N losses_i
     $$
+
+    Attributes:
+        should_persist(bool): Whether to persist the loss aggregator when saving.
+            Those loss aggregators with parameters and/or buffers should be persisted.
     """
+    should_persist: ClassVar[bool] = False
 
     def __init__(self) -> None:
         self.step = 0

ppsci/solver/solver.py

@@ -334,6 +334,7 @@ def __init__(
                 self.scaler,
                 self.equation,
                 self.ema_model,
+                self.loss_aggregator,
             )
             if isinstance(loaded_metric, dict):
                 self.best_metric.update(loaded_metric)
@@ -567,6 +568,7 @@ def train(self) -> None:
                         self.output_dir,
                         "best_model",
                         self.equation,
+                        aggregator=self.loss_aggregator,
                     )
                 logger.info(
                     f"[Eval][Epoch {epoch_id}]"
@@ -633,6 +635,7 @@ def train(self) -> None:
                     f"epoch_{epoch_id}",
                     self.equation,
                     ema_model=self.ema_model,
+                    aggregator=self.loss_aggregator,
                 )
 
             # save the latest model for convenient resume training
@@ -646,6 +649,7 @@ def train(self) -> None:
                 self.equation,
                 print_log=(epoch_id == start_epoch),
                 ema_model=self.ema_model,
+                aggregator=self.loss_aggregator,
             )
 
     def finetune(self, pretrained_model_path: str) -> None:

ppsci/utils/save_load.py

@@ -31,6 +31,7 @@
     from paddle import optimizer
 
     from ppsci import equation
+    from ppsci.loss import mtl
     from ppsci.utils import ema
 
 
@@ -42,7 +43,10 @@
 
 
 def _load_pretrain_from_path(
-    path: str, model: nn.Layer, equation: Optional[Dict[str, equation.PDE]] = None
+    path: str,
+    model: nn.Layer,
+    equation: Optional[Dict[str, equation.PDE]] = None,
+    loss_aggregator: Optional[mtl.LossAggregator] = None,
 ):
     """Load pretrained model from given path.
 
@@ -77,9 +81,26 @@ def _load_pretrain_from_path(
                 f"Finish loading pretrained equation parameters from: {path}.pdeqn"
             )
 
+    if loss_aggregator is not None:
+        if not os.path.exists(f"{path}.pdagg"):
+            if loss_aggregator.should_persist:
+                logger.warning(
+                    f"Given loss_aggregator({type(loss_aggregator)}) has persistable"
+                    f"parameters or buffers, but {path}.pdagg not found."
+                )
+        else:
+            aggregator_dict = paddle.load(f"{path}.pdagg")
+            loss_aggregator.set_state_dict(aggregator_dict)
+            logger.message(
+                f"Finish loading pretrained equation parameters from: {path}.pdagg"
+            )
+
 
 def load_pretrain(
-    model: nn.Layer, path: str, equation: Optional[Dict[str, equation.PDE]] = None
+    model: nn.Layer,
+    path: str,
+    equation: Optional[Dict[str, equation.PDE]] = None,
+    loss_aggregator: Optional[mtl.LossAggregator] = None,
 ):
     """
     Load pretrained model from given path or url.
@@ -121,7 +142,7 @@ def is_url_accessible(url: str):
     # remove ".pdparams" in suffix of path for convenient
     if path.endswith(".pdparams"):
         path = path[:-9]
-    _load_pretrain_from_path(path, model, equation)
+    _load_pretrain_from_path(path, model, equation, loss_aggregator)
 
 
 def load_checkpoint(
@@ -131,6 +152,7 @@ def load_checkpoint(
     grad_scaler: Optional[amp.GradScaler] = None,
     equation: Optional[Dict[str, equation.PDE]] = None,
     ema_model: Optional[ema.AveragedModel] = None,
+    aggregator: Optional[mtl.LossAggregator] = None,
 ) -> Dict[str, Any]:
     """Load from checkpoint.
 
@@ -141,6 +163,7 @@ def load_checkpoint(
         grad_scaler (Optional[amp.GradScaler]): GradScaler for AMP. Defaults to None.
         equation (Optional[Dict[str, equation.PDE]]): Equations. Defaults to None.
         ema_model: Optional[ema.AveragedModel]: Average model. Defaults to None.
+        aggregator: Optional[mtl.LossAggregator]: Loss aggregator. Defaults to None.
 
     Returns:
         Dict[str, Any]: Loaded metric information.
@@ -189,6 +212,10 @@ def load_checkpoint(
         avg_param_dict = paddle.load(f"{path}_ema.pdparams")
         ema_model.set_state_dict(avg_param_dict)
 
+    if aggregator is not None:
+        aggregator_dict = paddle.load(f"{path}.pdagg")
+        aggregator.set_state_dict(aggregator_dict)
+
     logger.message(f"Finish loading checkpoint from {path}")
     return metric_dict
 
@@ -203,6 +230,7 @@ def save_checkpoint(
     equation: Optional[Dict[str, equation.PDE]] = None,
     print_log: bool = True,
     ema_model: Optional[ema.AveragedModel] = None,
+    aggregator: Optional[mtl.LossAggregator] = None,
 ):
     """
     Save checkpoint, including model params, optimizer params, metric information.
@@ -219,6 +247,7 @@ def save_checkpoint(
             keeping log tidy without duplicate 'Finish saving checkpoint ...' log strings.
             Defaults to True.
         ema_model: Optional[ema.AveragedModel]: Average model. Defaults to None.
+        aggregator: Optional[mtl.LossAggregator]: Loss aggregator. Defaults to None.
 
     Examples:
         >>> import ppsci
@@ -258,6 +287,9 @@ def save_checkpoint(
     if ema_model:
         paddle.save(ema_model.state_dict(), f"{ckpt_path}_ema.pdparams")
 
+    if aggregator and aggregator.should_persist:
+        paddle.save(aggregator.state_dict(), f"{ckpt_path}.pdagg")
+
     if print_log:
         log_str = f"Finish saving checkpoint to: {ckpt_path}"
         if prefix == "latest":