refine save_load (#1158)

Author: HydrogenSulfate

ppsci/utils/save_load.py

@@ -160,9 +160,11 @@ def load_checkpoint(
         raise FileNotFoundError(f"{path}.scaler not exist.")
 
     # load state dict
-    param_dict = paddle.load(f"{path}.pdparams")
+    model_dict = paddle.load(f"{path}.pdparams")
     optim_dict = paddle.load(f"{path}.pdopt")
-    metric_dict = paddle.load(f"{path}.pdstates")
+    metric_dict = {}
+    if os.path.exists(f"{path}.pdstates"):
+        metric_dict = paddle.load(f"{path}.pdstates")
     if grad_scaler is not None:
         scaler_dict = paddle.load(f"{path}.pdscaler")
     if equation is not None:
@@ -172,9 +174,9 @@ def load_checkpoint(
         else:
             equation_dict = paddle.load(f"{path}.pdeqn")
 
-    # set state dict
+    # set model state dict
     logger.message(f"* Loading model checkpoint from {path}.pdparams")
-    missing_keys, unexpected_keys = model.set_state_dict(param_dict)
+    missing_keys, unexpected_keys = model.set_state_dict(model_dict)
     if missing_keys:
         logger.warning(
             f"There are missing keys when loading checkpoint: {missing_keys}, "
@@ -186,20 +188,23 @@ def load_checkpoint(
             "and corresponding weights will be ignored."
         )
 
+    # set optimizer state dict
     logger.message(f"* Loading optimizer checkpoint from {path}.pdopt")
     optimizer.set_state_dict(optim_dict)
+
     if grad_scaler is not None:
         logger.message(f"* Loading grad scaler checkpoint from {path}.pdscaler")
         grad_scaler.load_state_dict(scaler_dict)
+
     if equation is not None and equation_dict is not None:
         logger.message(f"* Loading equation checkpoint from {path}.pdeqn")
         for name, _equation in equation.items():
             _equation.set_state_dict(equation_dict[name])
 
-    if ema_model:
+    if ema_model is not None:
         logger.message(f"* Loading EMA checkpoint from {path}_ema.pdparams")
-        avg_param_dict = paddle.load(f"{path}_ema.pdparams")
-        ema_model.set_state_dict(avg_param_dict)
+        avg_model_dict = paddle.load(f"{path}_ema.pdparams")
+        ema_model.set_state_dict(avg_model_dict)
 
     if aggregator is not None and aggregator.should_persist:
         logger.message(f"* Loading loss aggregator checkpoint from {path}.pdagg")
@@ -213,7 +218,7 @@ def load_checkpoint(
 def save_checkpoint(
     model: nn.Layer,
     optimizer: Optional[optimizer.Optimizer],
-    metric: Dict[str, float],
+    metric: Optional[Dict[str, float]] = None,
     grad_scaler: Optional[amp.GradScaler] = None,
     output_dir: Optional[str] = None,
     prefix: str = "model",
@@ -228,7 +233,7 @@ def save_checkpoint(
     Args:
         model (nn.Layer): Model with parameters.
         optimizer (Optional[optimizer.Optimizer]): Optimizer for model.
-        metric (Dict[str, float]): Metric information, such as {"RMSE": 0.1, "MAE": 0.2}.
+        metric (Optional[Dict[str, float]]): Metric information, such as {"RMSE": 0.1, "MAE": 0.2}. Defaults to None.
         grad_scaler (Optional[amp.GradScaler]): GradScaler for AMP. Defaults to None.
         output_dir (Optional[str]): Directory for checkpoint storage.
         prefix (str, optional): Prefix for storage. Defaults to "model".
@@ -259,11 +264,16 @@ def save_checkpoint(
     os.makedirs(ckpt_dir, exist_ok=True)
 
     paddle.save(model.state_dict(), f"{ckpt_path}.pdparams")
-    if optimizer:
+
+    if optimizer is not None:
         paddle.save(optimizer.state_dict(), f"{ckpt_path}.pdopt")
-    paddle.save(metric, f"{ckpt_path}.pdstates")
+
+    if metric is not None and len(metric) > 0:
+        paddle.save(metric, f"{ckpt_path}.pdstates")
+
     if grad_scaler is not None:
         paddle.save(grad_scaler.state_dict(), f"{ckpt_path}.pdscaler")
+
     if equation is not None:
         num_learnable_params = sum(
             [len(eq.learnable_parameters) for eq in equation.values()]
@@ -274,10 +284,10 @@ def save_checkpoint(
                 f"{ckpt_path}.pdeqn",
             )
 
-    if ema_model:
+    if ema_model is not None:
         paddle.save(ema_model.state_dict(), f"{ckpt_path}_ema.pdparams")
 
-    if aggregator and aggregator.should_persist:
+    if aggregator is not None and aggregator.should_persist:
         paddle.save(aggregator.state_dict(), f"{ckpt_path}.pdagg")
 
     if print_log: