clarena.pipelines.mtl_eval

  1r"""The submodule in `pipelines` for multi-task learning evaluation."""
  2
  3__all__ = ["MTLEvaluation"]
  4
  5import logging
  6
  7import hydra
  8import lightning as L
  9import torch
 10from lightning import Callback, Trainer
 11from lightning.pytorch.loggers import Logger
 12from omegaconf import DictConfig, ListConfig
 13
 14from clarena.mtl_datasets import MTLDataset
 15
 16# always get logger for built-in logging in each module
 17pylogger = logging.getLogger(__name__)
 18
 19
 20class MTLEvaluation:
 21    r"""The base class for multi-task learning evaluation."""
 22
 23    def __init__(self, cfg: DictConfig) -> None:
 24        r"""
 25        **Args:**
 26        - **cfg** (`DictConfig`): the config dict for the multi-task learning evaluation.
 27        """
 28        self.cfg: DictConfig = cfg
 29        r"""The complete config dict."""
 30
 31        MTLEvaluation.sanity_check(self)
 32
 33        # required config fields
 34        self.eval_tasks: list[int] = (
 35            cfg.eval_tasks
 36            if isinstance(cfg.eval_tasks, ListConfig)
 37            else list(range(1, cfg.eval_tasks + 1))
 38        )
 39        r"""The list of task IDs to evaluate."""
 40        self.global_seed: int = cfg.global_seed
 41        r"""The global seed for the entire experiment."""
 42        self.output_dir: str = cfg.output_dir
 43        r"""The folder for storing the experiment results."""
 44        self.model_path: str = cfg.model_path
 45        r"""The file path of the model to evaluate."""
 46
 47        # components
 48        self.mtl_dataset: MTLDataset
 49        r"""MTL dataset object."""
 50        self.lightning_loggers: list[Logger]
 51        r"""Lightning logger objects."""
 52        self.callbacks: list[Callback]
 53        r"""Callback objects."""
 54        self.trainer: Trainer
 55        r"""Trainer object."""
 56
 57    def sanity_check(self) -> None:
 58        r"""Sanity check for config."""
 59
 60        # check required config fields
 61        required_config_fields = [
 62            "pipeline",
 63            "model_path",
 64            "eval_tasks",
 65            "global_seed",
 66            "trainer",
 67            "metrics",
 68            "callbacks",
 69            "output_dir",
 70            # "hydra" is excluded as it doesn't appear
 71            "misc",
 72        ]
 73        for field in required_config_fields:
 74            if not self.cfg.get(field):
 75                raise KeyError(
 76                    f"Field `{field}` is required in the experiment index config."
 77                )
 78
 79        # get dataset number of tasks
 80        if self.cfg.mtl_dataset._target_ == "clarena.mtl_datasets.MTLDatasetFromCL":
 81            cl_dataset_cfg = self.cfg.mtl_dataset.get("cl_dataset")
 82            if cl_dataset_cfg.get("num_tasks"):
 83                num_tasks = cl_dataset_cfg.get("num_tasks")
 84            elif cl_dataset_cfg.get("class_split"):
 85                num_tasks = len(cl_dataset_cfg.class_split)
 86            elif cl_dataset_cfg.get("datasets"):
 87                num_tasks = len(cl_dataset_cfg.datasets)
 88            else:
 89                raise KeyError(
 90                    "`num_tasks` is required in cl_dataset config under mtl_dataset config. Please specify `num_tasks` (for `CLPermutedDataset`) or `class_split` (for `CLSplitDataset`) or `datasets` (for `CLCombinedDataset`) in cl_dataset config."
 91                )
 92        else:
 93            if self.cfg.mtl_dataset.get("num_tasks"):
 94                num_tasks = self.cfg.mtl_dataset.num_tasks
 95            else:
 96                raise KeyError(
 97                    "`num_tasks` is required in mtl_dataset config. Please specify `num_tasks` in mtl_dataset config."
 98                )
 99
100        # check eval_tasks
101        eval_tasks = self.cfg.eval_tasks
102        if isinstance(eval_tasks, list):
103            if len(eval_tasks) < 1:
104                raise ValueError("`eval_tasks` must contain at least one task.")
105            if any(t < 1 or t > num_tasks for t in eval_tasks):
106                raise ValueError(
107                    f"All task IDs in `eval_tasks` must be between 1 and {num_tasks}."
108                )
109        elif isinstance(eval_tasks, int):
110            if eval_tasks < 0 or eval_tasks > num_tasks:
111                raise ValueError(
112                    f"`eval_tasks` as integer must be between 0 and {num_tasks}."
113                )
114        else:
115            raise TypeError(
116                "`eval_tasks` must be either a list of integers or an integer."
117            )
118
119    def instantiate_mtl_dataset(
120        self,
121        mtl_dataset_cfg: DictConfig,
122    ) -> None:
123        r"""Instantiate the MTL dataset object from `mtl_dataset_cfg`."""
124        pylogger.debug(
125            "Instantiating MTL dataset <%s> (clarena.mtl_datasets.MTLDataset) ...",
126            mtl_dataset_cfg.get("_target_"),
127        )
128        self.mtl_dataset = hydra.utils.instantiate(mtl_dataset_cfg)
129        pylogger.debug(
130            "MTL dataset <%s> (clarena.mtl_datasets.MTLDataset) instantiated! ",
131            mtl_dataset_cfg.get("_target_"),
132        )
133
134    def instantiate_callbacks(
135        self, metrics_cfg: DictConfig, callbacks_cfg: DictConfig
136    ) -> None:
137        r"""Instantiate the list of callbacks objects from `metrics_cfg` and `callbacks_cfg`. Note that `metrics_cfg` is a list of metric callbacks and `callbacks_cfg` is a list of callbacks other the metric callbacks. The instantiated callbacks contain both metric callbacks and other callbacks."""
138        pylogger.debug("Instantiating callbacks (lightning.Callback)...")
139
140        # instantiate metric callbacks
141        metric_callbacks = [
142            hydra.utils.instantiate(callback) for callback in metrics_cfg
143        ]
144
145        # instantiate other callbacks
146        other_callbacks = [
147            hydra.utils.instantiate(callback) for callback in callbacks_cfg
148        ]
149
150        # add metric callbacks to the list of callbacks
151        self.callbacks = metric_callbacks + other_callbacks
152        pylogger.debug("Callbacks (lightning.Callback) instantiated!")
153
154    def instantiate_trainer(
155        self,
156        trainer_cfg: DictConfig,
157        callbacks: list[Callback],
158    ) -> None:
159        r"""Instantiate the trainer object from `trainer_cfg` and `callbacks`."""
160
161        pylogger.debug("Instantiating trainer (lightning.Trainer)...")
162        self.trainer = hydra.utils.instantiate(trainer_cfg, callbacks=callbacks)
163        pylogger.debug("Trainer (lightning.Trainer) instantiated!")
164
165    def set_global_seed(self, global_seed: int) -> None:
166        r"""Set the `global_seed` for the entire evaluation."""
167        L.seed_everything(self.global_seed, workers=True)
168        pylogger.debug("Global seed is set as %d.", global_seed)
169
170    def run(self) -> None:
171        r"""The main method to run the multi-task learning evaluation."""
172
173        self.set_global_seed(self.global_seed)
174
175        # load the model from file
176        model = torch.load(self.model_path)
177
178        self.instantiate_mtl_dataset(mtl_dataset_cfg=self.cfg.mtl_dataset)
179        self.instantiate_callbacks(
180            metrics_cfg=self.cfg.metrics, callbacks_cfg=self.cfg.callbacks
181        )
182        self.instantiate_trainer(
183            trainer_cfg=self.cfg.trainer,
184            callbacks=self.callbacks,
185        )  # trainer should be instantiated after callbacks
186
187        # setup tasks for dataset
188        self.mtl_dataset.setup_tasks_eval(eval_tasks=self.eval_tasks)
189
190        # evaluation skipping training and validation
191        self.trainer.test(
192            model=model,
193            datamodule=self.mtl_dataset,
194        )