correct LBFGS code

Author: HydrogenSulfate

ppsci/solver/eval.py

@@ -66,7 +66,7 @@ def eval_func(solver, epoch_id: int, log_freq: int) -> float:
                 evaluator.add_target_expr(output_formula, output_name)
 
             # forward
-            with solver._autocast_context_manager():
+            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)

ppsci/solver/solver.py

@@ -27,9 +27,8 @@
 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 import optimizer as optim
 from paddle.distributed import fleet
 from typing_extensions import Literal
 
@@ -100,8 +99,8 @@ def __init__(
         model: nn.Layer,
         constraint: Optional[Dict[str, ppsci.constraint.Constraint]] = None,
         output_dir: str = "./output/",
-        optimizer: Optional[optimizer.Optimizer] = None,
-        lr_scheduler: Optional[optimizer.lr.LRScheduler] = None,
+        optimizer: Optional[optim.Optimizer] = None,
+        lr_scheduler: Optional[optim.lr.LRScheduler] = None,
         epochs: int = 5,
         iters_per_epoch: int = 20,
         update_freq: int = 1,
@@ -205,10 +204,10 @@ def __init__(
                 self.best_metric.update(loaded_metric)
 
         # choosing an appropriate training function for different optimizers
-        if not isinstance(self.optimizer, incubate.optimizer.LBFGS):
-            self.train_epoch_func = ppsci.solver.train.train_epoch_func
-        else:
+        if isinstance(self.optimizer, optim.LBFGS):
             self.train_epoch_func = ppsci.solver.train.train_LBFGS_epoch_func
+        else:
+            self.train_epoch_func = ppsci.solver.train.train_epoch_func
 
         # decorate model(s) and optimizer(s) for AMP
         if self.use_amp:
@@ -471,7 +470,7 @@ def predict(
                 batch_input_dict[key].stop_gradient = False
 
             # forward
-            with self._autocast_context_manager():
+            with self.autocast_context_manager():
                 batch_output_dict = self.model(batch_input_dict)
 
             # collect batch data
@@ -501,7 +500,7 @@ def export(self):
         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:
+    def autocast_context_manager(self) -> contextlib.AbstractContextManager:
         """Autocast context manager for Auto Mix Precision.
 
         Returns:

ppsci/solver/train.py

@@ -59,7 +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
-            with solver._autocast_context_manager():
+            with solver.autocast_context_manager():
                 output_dict = evaluator(input_dict)
                 constraint_loss = _constraint.loss(output_dict, label_dict, weight_dict)
                 total_loss += constraint_loss
@@ -114,16 +114,16 @@ def train_LBFGS_epoch_func(solver, epoch_id: int, log_freq: int):
     batch_tic = time.perf_counter()
 
     for iter_id in range(1, solver.iters_per_epoch + 1):
-        reader_cost = 0
-        batch_cost = 0
         loss_dict = misc.Prettydefaultdict(float)
         loss_dict["loss"] = 0.0
+        total_batch_size = []
+        reader_cost = 0
+        batch_cost = 0
+        reader_tic = time.perf_counter()
+
         input_dict_list = []
         label_dict_list = []
         weight_dict_list = []
-        batch_cost = 0
-        total_batch_size = []
-        reader_tic = time.perf_counter()
         for _, _constraint in solver.constraint.items():
             input_dict, label_dict, weight_dict = next(_constraint.data_iter)
             reader_cost += time.perf_counter() - reader_tic
@@ -133,10 +133,9 @@ def train_LBFGS_epoch_func(solver, epoch_id: int, log_freq: int):
             label_dict_list.append(label_dict)
             weight_dict_list.append(weight_dict)
             total_batch_size.append(next(iter(input_dict.values())).shape[0])
-        total_batch_size = sum(total_batch_size)
 
         def closure():
-            """Closure function for LBFGS optimizer.
+            """Forward-backward closure function for LBFGS optimizer.
 
             Returns:
                 Tensor: Computed loss.
@@ -149,21 +148,25 @@ def closure():
                 for output_name, output_formula in _constraint.output_expr.items():
                     evaluator.add_target_expr(output_formula, output_name)
 
-                # forward for every constraint
-                output_dict_i = evaluator(input_dict_list[i])
-                constraint_loss = _constraint.loss(
-                    output_dict_i, label_dict_list[i], weight_dict_list[i]
-                )
-                total_loss += constraint_loss
+                # forward for every batched data dict
+                with solver.autocast_context_manager():
+                    output_dict_i = evaluator(input_dict_list[i])
+                    constraint_loss = _constraint.loss(
+                        output_dict_i, label_dict_list[i], weight_dict_list[i]
+                    )
+                    total_loss += constraint_loss
 
                 loss_dict[_constraint.name] += float(constraint_loss)
 
             total_loss.backward()
+            loss_dict["loss"] = float(total_loss)
 
             return total_loss
 
+        reader_tic = time.perf_counter()
+
         solver.optimizer.step(closure)
-        if not getattr(solver.lr_scheduler, "by_epoch", False):
+        if solver.lr_scheduler is not None and not solver.lr_scheduler.by_epoch:
             solver.lr_scheduler.step()
 
         batch_cost += time.perf_counter() - batch_tic

ppsci/solver/visu.py

@@ -63,7 +63,7 @@ def visualize_func(solver, epoch_id: int):
                 evaluator.add_target_expr(output_expr, output_key)
 
             # forward
-            with solver._autocast_context_manager():
+            with solver.autocast_context_manager():
                 batch_output_dict = evaluator(batch_input_dict)
 
             # collect batch data