clarena.callbacks.pylogger

  1r"""
  2The submodule in `callbacks` for Pylogger callbacks.
  3"""
  4
  5__all__ = ["CLPylogger", "CULPylogger", "MTLPylogger", "STLPylogger"]
  6
  7import logging
  8
  9from lightning import Callback, Trainer
 10
 11from clarena.cl_algorithms import CLAlgorithm
 12from clarena.mtl_algorithms import MTLAlgorithm
 13from clarena.stl_algorithms import STLAlgorithm
 14
 15# always get logger for built-in logging in each module
 16pylogger = logging.getLogger(__name__)
 17
 18
 19class CLPylogger(Callback):
 20    r"""Provides additional logging messages for during continual learning progress.
 21
 22    Put logging messages here if you don't want to mess up the `CLAlgorithm` (`LightningModule`) with a huge amount of logging codes.
 23    """
 24
 25    def on_fit_start(self, trainer: Trainer, pl_module: CLAlgorithm) -> None:
 26        r"""Log messages for the start of training task."""
 27        pylogger.info("Start training continual learning task %s!", pl_module.task_id)
 28
 29    def on_train_end(self, trainer: Trainer, pl_module: CLAlgorithm) -> None:
 30        r"""Log messages for the end of training task."""
 31        pylogger.info("Finish training continual learning task %s!", pl_module.task_id)
 32
 33    def on_test_start(self, trainer: Trainer, pl_module: CLAlgorithm) -> None:
 34        r"""Log messages for the start of testing task."""
 35        pylogger.info(
 36            "Start testing continual learning task %s on all previous and current tasks!",
 37            pl_module.task_id,
 38        )
 39
 40    def on_test_end(self, trainer: Trainer, pl_module: CLAlgorithm) -> None:
 41        r"""Log messages for the end of testing task."""
 42        pylogger.info(
 43            "Finish testing continual learning task %s on all previous and current tasks!",
 44            pl_module.task_id,
 45        )
 46
 47
 48class CULPylogger(Callback):
 49    r"""Provides additional logging messages for during continual learning progress.
 50
 51    Put logging messages here if you don't want to mess up the `CLAlgorithm` (`LightningModule`) with a huge amount of logging codes.
 52    """
 53
 54    def on_fit_start(self, trainer: Trainer, pl_module: CLAlgorithm) -> None:
 55        r"""Log messages for the start of training task."""
 56        pylogger.info("Start training continual learning task %s!", pl_module.task_id)
 57
 58    def on_train_end(self, trainer: Trainer, pl_module: CLAlgorithm) -> None:
 59        r"""Log messages for the end of training task."""
 60        pylogger.info("Finish training continual learning task %s!", pl_module.task_id)
 61
 62    def on_test_start(self, trainer: Trainer, pl_module: CLAlgorithm) -> None:
 63        r"""Log messages for the start of testing task."""
 64        pylogger.info(
 65            "Start testing continual learning task %s on all previous and current tasks!",
 66            pl_module.task_id,
 67        )
 68
 69    def on_test_end(self, trainer: Trainer, pl_module: CLAlgorithm) -> None:
 70        r"""Log messages for the end of testing task."""
 71        pylogger.info(
 72            "Finish testing continual learning task %s on all previous and current tasks!",
 73            pl_module.task_id,
 74        )
 75
 76
 77class MTLPylogger(Callback):
 78    r"""Pylogger Callback provides additional logging for during multi-task learning progress.
 79
 80    Put logging messages here if you don't want to mess up the `MTLAlgorithm` (`LightningModule`) with a huge amount of logging codes.
 81    """
 82
 83    def on_fit_start(self, trainer: Trainer, pl_module: MTLAlgorithm) -> None:
 84        r"""Log messages for the start of training."""
 85        pylogger.info("Start multi-task training!")
 86
 87    def on_train_end(self, trainer: Trainer, pl_module: MTLAlgorithm) -> None:
 88        r"""Log messages for the end of training."""
 89        pylogger.info("Finish multi-task training!")
 90
 91    def on_test_start(self, trainer: Trainer, pl_module: MTLAlgorithm) -> None:
 92        r"""Log messages for the start of testing."""
 93        pylogger.info("Start testing multi-task testing!")
 94
 95    def on_test_end(self, trainer: Trainer, pl_module: MTLAlgorithm) -> None:
 96        r"""Log messages for the end of testing."""
 97        pylogger.info("Finish testing multi-task testing!")
 98
 99
100class STLPylogger(Callback):
101    r"""Pylogger Callback provides additional logging for during single-task learning progress.
102
103    Put logging messages here if you don't want to mess up the `STLAlgorithm` (`LightningModule`) with a huge amount of logging codes.
104    """
105
106    def on_fit_start(self, trainer: Trainer, pl_module: STLAlgorithm) -> None:
107        r"""Log messages for the start of training."""
108        pylogger.info("Start single-task training!")
109
110    def on_train_end(self, trainer: Trainer, pl_module: STLAlgorithm) -> None:
111        r"""Log messages for the end of training."""
112        pylogger.info("Finish single-task training!")
113
114    def on_test_start(self, trainer: Trainer, pl_module: STLAlgorithm) -> None:
115        r"""Log messages for the start of testing."""
116        pylogger.info("Start single-task testing!")
117
118    def on_test_end(self, trainer: Trainer, pl_module: STLAlgorithm) -> None:
119        r"""Log messages for the end of testing."""
120        pylogger.info("Finish single-task testing!")