[Fix] Fix aggregator error when load ckpt (#1038)

* fix aggregator error when load ckpt

* add unitest

Author: HydrogenSulfate

ppsci/solver/solver.py

@@ -332,6 +332,15 @@ def __init__(
             "metric": float("inf"),
             "epoch": 0,
         }
+
+        # use loss aggregator, use Sum if None
+        if isinstance(loss_aggregator, (mtl.AGDA, mtl.PCGrad)) and self.use_amp:
+            raise ValueError(
+                "Auto Mix Precision do not support AGDA, PCGrad loss aggregator yet, "
+                "please set use_amp=False."
+            )
+        self.loss_aggregator = loss_aggregator or mtl.Sum()
+
         # load model checkpoint, usually used for resume training
         if not cfg:
             self.checkpoint_path = checkpoint_path
@@ -478,14 +487,6 @@ def dist_wrapper(model: nn.Layer) -> paddle.DataParallel:
         jit.enable_to_static(to_static)
         logger.message(f"Set to_static={to_static} for computational optimization.")
 
-        # use loss aggregator, use Sum if None
-        if isinstance(loss_aggregator, (mtl.AGDA, mtl.PCGrad)) and self.use_amp:
-            raise ValueError(
-                "Auto Mix Precision do not support AGDA, PCGrad loss aggregator yet, "
-                "please set use_amp=False."
-            )
-        self.loss_aggregator = loss_aggregator or mtl.Sum()
-
         # convert sympy to callable object if exist
         extra_parameters = []
         if self.equation:

ppsci/utils/save_load.py

@@ -196,7 +196,7 @@ 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:
+    if aggregator is not None and aggregator.should_persist:
         aggregator_dict = paddle.load(f"{path}.pdagg")
         aggregator.set_state_dict(aggregator_dict)
 

test/loss/aggregator.py

@@ -0,0 +1,98 @@
+import pytest
+
+import ppsci
+from ppsci import arch
+from ppsci.loss import mtl
+
+__all__ = []
+
+
+class AggregatorTest:
+    def __init__(self):
+        self.model = arch.MLP(
+            ("x", "y"),
+            ("u", "v"),
+            3,
+            16,
+        )
+
+    def _check_agg_state_dict(self, agg):
+        model_state = self.model.state_dict()
+        agg_state = agg.state_dict()
+        for k in agg_state:
+            assert k not in model_state
+
+    def test_AGDA(self):
+        aggregator = mtl.AGDA(self.model)
+        assert aggregator.should_persist is False
+
+    def test_GradNorm(self):
+        aggregator = mtl.GradNorm(self.model)
+        assert aggregator.should_persist is True
+        self._check_agg_state_dict(aggregator)
+
+    def test_LossAggregator(self):
+        aggregator = mtl.AGDA(self.model)
+        assert aggregator.should_persist is False
+
+    def test_PCGrad(self):
+        aggregator = mtl.PCGrad(self.model)
+        assert aggregator.should_persist is False
+
+    def test_Relobralo(self):
+        aggregator = mtl.Relobralo(self.model)
+        assert aggregator.should_persist is True
+        self._check_agg_state_dict(aggregator)
+
+    def test_Sum(self):
+        aggregator = mtl.Sum(self.model)
+        assert aggregator.should_persist is False
+
+    def test_NTK(self):
+        aggregator = mtl.NTK(self.model)
+        assert aggregator.should_persist is True
+        self._check_agg_state_dict(aggregator)
+
+    def test_restore_aggregator(self):
+        model = ppsci.arch.MLP(
+            ["x", "y"],
+            ["u"],
+            2,
+            16,
+        )
+        opt = ppsci.optimizer.Adam(1e-3)(model)
+        equation = ppsci.equation.Laplace(2)
+        geom = ppsci.geometry.Rectangle([0, 0], [1, 1])
+        BC = ppsci.constraint.BoundaryConstraint(
+            equation.equations,
+            {"laplace": 0.0},
+            geom,
+            {
+                "dataset": "IterableNamedArrayDataset",
+                "iters_per_epoch": 10,
+                "batch_size": 16,
+            },
+            loss=ppsci.loss.MSELoss(),
+        )
+        solver = ppsci.solver.Solver(
+            model,
+            {"bound": BC},
+            optimizer=opt,
+            output_dir="./tmp",
+            iters_per_epoch=10,
+            epochs=2,
+        )
+        solver.train()
+        solver = ppsci.solver.Solver(
+            model,
+            {"bound": BC},
+            optimizer=opt,
+            output_dir="./tmp",
+            iters_per_epoch=10,
+            epochs=2,
+            checkpoint_path="./tmp/checkpoints/latest",
+        )
+
+
+if __name__ == "__main__":
+    pytest.main()