clarena.callbacks.pylogger

 1r"""
 2The submodule in `callbacks` for `PyloggerCallback`.
 3"""
 4
 5__all__ = ["PyloggerCallback"]
 6
 7import logging
 8
 9from lightning import Callback, Trainer
10
11from clarena.cl_algorithms import CLAlgorithm
12
13# always get logger for built-in logging in each module
14pylogger = logging.getLogger(__name__)
15
16
17class PyloggerCallback(Callback):
18    r"""Pylogger Callback provides additional logging for during continual learning progress.
19
20    Put logging messages here if you don't want to mess up the `CLAlgorithm` (`LightningModule`) with a huge amount of logging codes.
21    """
22
23    def on_fit_start(self, trainer: Trainer, pl_module: CLAlgorithm) -> None:
24        r"""Log messages for the start of training task."""
25        pylogger.info("Start training task %s!", pl_module.task_id)
26
27    def on_train_end(self, trainer: Trainer, pl_module: CLAlgorithm) -> None:
28        r"""Log messages for the end of training task."""
29        pylogger.info("Finish training task %s!", pl_module.task_id)
30
31    def on_test_start(self, trainer: Trainer, pl_module: CLAlgorithm) -> None:
32        r"""Log messages for the start of testing task."""
33        pylogger.info(
34            "Start testing task %s on all previous and current tasks!",
35            pl_module.task_id,
36        )
37
38    def on_test_end(self, trainer: Trainer, pl_module: CLAlgorithm) -> None:
39        r"""Log messages for the end of testing task."""
40        pylogger.info(
41            "Finish testing task %s on all previous and current tasks!",
42            pl_module.task_id,
43        )