support callback on epoch/iter begin&end and turorial (PaddlePaddle#1153)

Author: HydrogenSulfate

docs/zh/user_guide.md

@@ -1079,6 +1079,53 @@ PaddleScience 内置了两种模型平均方法：[Stochastic weight averaging(S
     3. 设置平均间隔为 1 个 epoch
     4. 设置平均的起始和终止 epoch 为 75 至 100
 
+### 2.7 回调(callback)注册与调用指南
+
+在深度学习模型的训练过程中，能够在特定的时机执行自定义逻辑是非常有用的。PaddleScience 的 `Solver` 类提供了一种相对灵活的机制，允许用户在**训练的不同阶段**注册和调用回调函数。
+
+具体地，我们提供了如下四种注册回调函数的接口：
+
+``` py
+Solver.register_callback_on_epoch_begin # 在每个 epoch 开始时调用
+Solver.register_callback_on_epoch_end # 在每个 epoch 结束时调用
+Solver.register_callback_on_iter_begin # 在每个 iteration 开始时调用
+Solver.register_callback_on_iter_end # 在每个 iteration 结束时调用
+```
+
+它们在训练过程中的调用时机如下示例所示：
+
+``` py hl_lines="3 6 8 10"
+for epoch_id in range(1, num_epochs + 1):
+    # train one epoch...
+    _invoke_callbacks_on_epoch_begin() # 此处按注册顺序, 自动调用通过 register_callback_on_epoch_begin 注册的回调函数
+
+    for iter_id in range(1, num_iters + 1)
+        _invoke_callbacks_on_iter_begin() # 此处按注册顺序, 自动调用通过 register_callback_on_iter_begin 注册的回调函数
+        # train one iteration...
+        _invoke_callbacks_on_iter_end() # 此处按注册顺序, 自动调用通过 register_callback_on_iter_end 注册的回调函数
+
+    _invoke_callbacks_on_epoch_end() # 此处按注册顺序, 自动调用通过 register_callback_on_epoch_end 注册的回调函数
+```
+
+以 `examples/fsi/viv.py` 为例，假设希望在训练时，每隔 100 个 epoch 打印出方程中的可学习参数 `k1`, `k2`，那么可以按照如下示例代码，添加回调函数：
+
+``` py hl_lines="11 12 13 14 15"
+# initialize solver
+solver = ppsci.solver.Solver(
+    model,
+    constraint,
+    optimizer=optimizer,
+    equation=equation,
+    validator=validator,
+    visualizer=visualizer,
+    cfg=cfg,
+)
+def show_learnable_params(slv):
+    if slv.global_step % 100 == 0:
+        ppsci.utils.logger.message(f"{equation['VIV'].k1.item():.5f}, {equation['VIV'].k2.item():.5f}")
+solver.register_callback_on_iter_begin(show_learnable_params)
+```
+
 ## 3. 使用 Nsight 进行性能分析
 
 Nsight是NVIDIA面向开发者提供的开发工具套件，能提供深入的跟踪、调试、评测和分析，以优化跨 NVIDIA GPU和CPU的复杂计算应用程序。详细文档可参考：[Nsight Systems Document](https://docs.nvidia.com/nsight-systems/index.html)

ppsci/solver/solver.py

@@ -552,6 +552,12 @@ def convert_expr(
         self.nvtx_flag: bool = os.getenv("NVTX", None) is not None
         self.forward_helper.nvtx_flag = self.nvtx_flag
 
+        # for callbacks
+        self.callbacks_on_epoch_begin: List[Callable[[Solver]]] = []
+        self.callbacks_on_epoch_end: List[Callable[[Solver]]] = []
+        self.callbacks_on_iter_begin: List[Callable[[Solver]]] = []
+        self.callbacks_on_iter_end: List[Callable[[Solver]]] = []
+
     def train(self) -> None:
         """Training."""
         self.global_step = self.best_metric["epoch"] * self.iters_per_epoch
@@ -569,7 +575,10 @@ def train(self) -> None:
             core.nvprof_enable_record_event()
 
         for epoch_id in range(start_epoch, self.epochs + 1):
+            self._invoke_callbacks_on_epoch_begin()  # [optional]
             self.train_epoch_func(self, epoch_id, self.log_freq)
+            self._invoke_callbacks_on_epoch_end()  # [optional]
+
             self.train_output_info.clear()
 
             # update average model if exist
@@ -1124,3 +1133,87 @@ def _parse_params_from_cfg(self, cfg: DictConfig):
             self.pretrained_model_path = cfg.EVAL.pretrained_model_path
         elif cfg.mode in ["export", "infer"]:
             self.pretrained_model_path = cfg.INFER.pretrained_model_path
+
+    def register_callback_on_epoch_begin(
+        self: Solver, callback_fn: Callable[[Solver]]
+    ) -> None:
+        """
+        Registers a callback function to be executed at the beginning of each training epoch.
+
+        Args:
+            callback_fn : Callable[[Solver]]
+                A function that takes a Solver instance as an argument. This function
+                will be called at the start of every epoch.
+        """
+        self.callbacks_on_epoch_begin.append(callback_fn)
+
+    def register_callback_on_epoch_end(
+        self: Solver, callback_fn: Callable[[Solver]]
+    ) -> None:
+        """
+        Registers a callback function to be executed at the end of each training epoch.
+
+        Args:
+            callback_fn : Callable[[Solver]]
+                A function that takes a Solver instance as an argument. This function
+                will be called at the end of every epoch.
+        """
+        self.callbacks_on_epoch_end.append(callback_fn)
+
+    def register_callback_on_iter_begin(
+        self: Solver, callback_fn: Callable[[Solver]]
+    ) -> None:
+        """
+        Registers a callback function to be executed at the beginning of each training iteration.
+
+        Args:
+            callback_fn : Callable[[Solver]]
+                A function that takes a Solver instance as an argument. This function
+                will be called at the start of every iteration.
+        """
+        self.callbacks_on_iter_begin.append(callback_fn)
+
+    def register_callback_on_iter_end(
+        self: Solver, callback_fn: Callable[[Solver]]
+    ) -> None:
+        """
+        Registers a callback function to be executed at the end of each training iteration.
+
+        Args:
+            callback_fn : Callable[[Solver]]
+                A function that takes a Solver instance as an argument. This function
+                will be called at the end of every iteration.
+
+        Returns:
+        -------
+        None
+        """
+        self.callbacks_on_iter_end.append(callback_fn)
+
+    def _invoke_callbacks_on_epoch_begin(self: Solver) -> None:
+        """
+        Invokes all registered callbacks at the beginning of an epoch.
+        """
+        for callback in self.callbacks_on_epoch_begin:
+            callback(self)
+
+    def _invoke_callbacks_on_epoch_end(self: Solver) -> None:
+        """
+        Invokes all registered callbacks at the end of an epoch.
+        """
+        for callback in self.callbacks_on_epoch_end:
+            callback(self)
+
+    def _invoke_callbacks_on_iter_begin(self: Solver) -> None:
+        """
+        Invokes all registered callbacks at the beginning of an iteration.
+        """
+        for callback in self.callbacks_on_iter_begin:
+            callback(self)
+
+    def _invoke_callbacks_on_iter_end(self: Solver) -> None:
+        """
+        Invokes all registered callbacks at the end of an iteration.
+        """
+        for callback in self.callbacks_on_iter_end:
+            callback(self)

ppsci/solver/train.py

@@ -66,6 +66,7 @@ def train_epoch_func(solver: "solver.Solver", epoch_id: int, log_freq: int):
     batch_tic = time.perf_counter()
 
     for iter_id in range(1, solver.iters_per_epoch + 1):
+        solver._invoke_callbacks_on_iter_begin()
         if solver.nvtx_flag:  # only for nsight analysis
             core.nvprof_nvtx_push(
                 f"Training iteration {solver.global_step + 1}"
@@ -212,6 +213,8 @@ def train_epoch_func(solver: "solver.Solver", epoch_id: int, log_freq: int):
                 core.nvprof_stop()
                 sys.exit(0)
 
+        solver._invoke_callbacks_on_iter_end()
+
 
 def train_LBFGS_epoch_func(solver: "solver.Solver", epoch_id: int, log_freq: int):
     """Train function for one epoch with L-BFGS optimizer.
@@ -226,6 +229,7 @@ def train_LBFGS_epoch_func(solver: "solver.Solver", epoch_id: int, log_freq: int
     batch_tic = time.perf_counter()
 
     for iter_id in range(1, solver.iters_per_epoch + 1):
+        solver._invoke_callbacks_on_iter_begin()
         loss_dict = misc.Prettydefaultdict(float)
         loss_dict["loss"] = 0.0
         total_batch_size = 0
@@ -317,3 +321,4 @@ def closure() -> paddle.Tensor:
             printer.log_train_info(solver, total_batch_size, epoch_id, iter_id)
 
         batch_tic = time.perf_counter()
+        solver._invoke_callbacks_on_iter_end()