clarena.cl_algorithms.independent

  1r"""
  2The submodule in `cl_algorithms` for Independent learning algorithm.
  3"""
  4
  5__all__ = ["Independent"]
  6
  7import inspect
  8import logging
  9from copy import deepcopy
 10from typing import Any
 11
 12from torch import Tensor
 13from torch.utils.data import DataLoader
 14
 15from clarena.backbones import CLBackbone
 16from clarena.cl_algorithms import Finetuning, UnlearnableCLAlgorithm
 17from clarena.heads import HeadDIL, HeadsCIL, HeadsTIL
 18
 19# always get logger for built-in logging in each module
 20pylogger = logging.getLogger(__name__)
 21
 22
 23class Independent(Finetuning):
 24    r"""Independent learning algorithm.
 25
 26    Another naive way for task-incremental learning aside from Finetuning.
 27    It assigns a new independent model for each task. This is a simple way
 28    to avoid catastrophic forgetting at the extreme cost of memory.
 29    It achieves the theoretical upper bound of performance in continual learning.
 30
 31    We implement Independent as a subclass of Finetuning algorithm, as
 32    Independent has the same `forward()`, `training_step()`,
 33    `validation_step()` and `test_step()` method as `Finetuning` class.
 34    """
 35
 36    def __init__(
 37        self,
 38        backbone: CLBackbone,
 39        heads: HeadsTIL | HeadsCIL | HeadDIL,
 40        non_algorithmic_hparams: dict[str, Any] = {},
 41        **kwargs,
 42    ) -> None:
 43        r"""Initialize the Independent algorithm with the network.
 44
 45        It has no additional hyperparameters.
 46
 47        **Args:**
 48        - **backbone** (`CLBackbone`): backbone network.
 49        - **heads** (`HeadsTIL`): output heads.
 50        - **non_algorithmic_hparams** (`dict[str, Any]`): non-algorithmic hyperparameters that are not related to the algorithm itself are passed to this `LightningModule` object from the config, such as optimizer and learning rate scheduler configurations.
 51        - **kwargs**: Reserved for multiple inheritance.
 52        """
 53        super().__init__(
 54            backbone=backbone,
 55            heads=heads,
 56            non_algorithmic_hparams=non_algorithmic_hparams,
 57            **kwargs,
 58        )
 59
 60        self.original_backbone_state_dict: dict = deepcopy(backbone.state_dict())
 61        r"""The original backbone state dict before training on any task. Used to initialize new independent backbones for new tasks."""
 62
 63        self.backbones: dict[int, CLBackbone] = {}
 64        r"""Independent backbones for each task.
 65        Keys are task IDs and values are the corresponding backbones.
 66        """
 67
 68        self.backbone_valid_task_ids: set[int] = set()
 69        r"""Task IDs that currently have valid trained backbones."""
 70
 71    def on_train_start(self):
 72        r"""At the start of training for current task `self.task_id`, load the original backbone state dict to create a new independent backbone for the current task."""
 73        self.backbone.load_state_dict(self.original_backbone_state_dict)
 74
 75    def on_train_end(self) -> None:
 76        r"""Save the trained independent backbone for `self.task_id`."""
 77        self.backbones[self.task_id] = deepcopy(self.backbone)
 78        self.backbone_valid_task_ids.add(self.task_id)
 79
 80    def test_step(
 81        self,
 82        batch: DataLoader,
 83        batch_idx: int,
 84        dataloader_idx: int = 0,
 85    ) -> dict[str, Tensor]:
 86        r"""Test step for Independent.
 87
 88        Tests all seen tasks indexed by `dataloader_idx`.
 89
 90        **Args:**
 91        - **batch** (`Any`): a batch of test data.
 92        - **dataloader_idx** (`int`): task ID being tested.
 93
 94        **Returns:**
 95        - **outputs** (`dict[str, Tensor]`): loss and accuracy.
 96        """
 97        test_task_id = self.get_test_task_id_from_dataloader_idx(dataloader_idx)
 98
 99        x, y = batch
100        backbone = self.backbones[test_task_id]
101        feature, _ = backbone(x, stage="test", task_id=test_task_id)
102        logits = self.heads(feature, test_task_id)
103
104        loss_cls = self.criterion(logits, y)
105        preds = logits.argmax(dim=1)
106        acc = (preds == y).float().mean()
107
108        return {
109            "loss_cls": loss_cls,
110            "acc": acc,
111            "preds": preds,
112        }
113
114
115class UnlearnableIndependent(UnlearnableCLAlgorithm, Independent):
116    r"""Unlearnable Independent learning algorithm.
117
118    Variant of Independent that supports unlearning requests and
119    permanent tasks.
120    """
121
122    def __init__(
123        self,
124        backbone: CLBackbone,
125        heads: HeadsTIL,
126        non_algorithmic_hparams: dict[str, Any] = {},
127        disable_unlearning: bool = False,
128    ) -> None:
129        r"""Initialize Unlearnable Independent.
130
131        **Args:**
132        - **backbone** (`CLBackbone`): backbone network.
133        - **heads** (`HeadsTIL`): output heads. Independent supports only TIL (Task-Incremental Learning).
134        - **disable_unlearning** (`bool`): disable unlearning or not.
135        """
136        super().__init__(
137            backbone=backbone,
138            heads=heads,
139            non_algorithmic_hparams=non_algorithmic_hparams,
140            disable_unlearning=disable_unlearning,
141        )
142
143    def aggregated_backbone_output(self, input: Tensor) -> Tensor:
144        r"""Aggregate backbone outputs across all valid tasks.
145
146        This feature is used for unlearning metrics such as Distribution Distance (DD).
147
148        **Args:**
149        - **input** (`Tensor`): input tensor.
150
151        **Returns:**
152        - **output** (`Tensor`): aggregated backbone feature.
153        """
154        feature = 0
155
156        for t in self.backbone_valid_task_ids:
157            backbone_t = self.backbones[t]
158            forward_params = inspect.signature(backbone_t.forward).parameters
159            supports_test_task_id = "test_task_id" in forward_params or any(
160                p.kind == inspect.Parameter.VAR_KEYWORD for p in forward_params.values()
161            )
162            if supports_test_task_id:
163                feature_t = backbone_t(input, stage="unlearning_test", test_task_id=t)[
164                    0
165                ]
166            else:
167                feature_t = backbone_t(input, stage="unlearning_test")[0]
168            feature += feature_t
169
170        feature = feature / len(self.backbone_valid_task_ids)
171        return feature
172
173
174# r"""
175# The submodule in `cl_algorithms` for Independent learning algorithm.
176# """
177
178# __all__ = ["Independent"]
179
180# import logging
181# from copy import deepcopy
182# from typing import Any
183
184# from torch import Tensor
185# from torch.optim import Optimizer
186# from torch.optim.lr_scheduler import LRScheduler
187# from torch.utils.data import DataLoader
188
189# from clarena.backbones import CLBackbone
190# from clarena.cl_algorithms import Finetuning, UnlearnableCLAlgorithm
191# from clarena.heads import HeadDIL, HeadsCIL, HeadsTIL
192
193# # always get logger for built-in logging in each module
194# pylogger = logging.getLogger(__name__)
195
196
197# class Independent(Finetuning):
198#     r"""Independent learning algorithm.
199
200#     Another naive way for task-incremental learning aside from Finetuning. It assigns a new independent model for each task. This is a simple way to avoid catastrophic forgetting at the extreme cost of memory. It achieves the theoretical upper bound of performance in continual learning.
201
202#     We implement Independent as a subclass of Finetuning algorithm, as Independent has the same `forward()`, `training_step()`, `validation_step()` and `test_step()` method as `Finetuning` class.
203#     """
204
205#     def __init__(
206#         self,
207#         backbone: CLBackbone,
208#         heads: HeadsTIL | HeadsCIL | HeadDIL,
209#         non_algorithmic_hparams: dict[str, Any] = {},
210#         **kwargs,
211#     ) -> None:
212#         r"""Initialize the Independent algorithm with the network. It has no additional hyperparameters.
213
214#         **Args:**
215#         - **backbone** (`CLBackbone`): backbone network.
216#         - **heads** (`HeadsTIL` | `HeadsCIL` | `HeadDIL`): output heads.
217#         - **non_algorithmic_hparams** (`dict[str, Any]`): non-algorithmic hyperparameters that are not related to the algorithm itself are passed to this `LightningModule` object from the config, such as optimizer and learning rate scheduler configurations. They are saved for Lightning APIs from `save_hyperparameters()` method. This is useful for the experiment configuration and reproducibility.
218#         - **kwargs**: Reserved for multiple inheritance.
219
220#         """
221#         super().__init__(
222#             backbone=backbone,
223#             heads=heads,
224#             non_algorithmic_hparams=non_algorithmic_hparams,
225#             **kwargs,
226#         )
227
228#         self.original_backbone_state_dict: dict = deepcopy(backbone.state_dict())
229#         r"""The original backbone state dict before training on any task. Used to initialize new independent backbones for new tasks."""
230
231#         self.backbones: dict[int, CLBackbone] = {}
232#         r"""The list of independent backbones for each task. Keys are task IDs and values are the corresponding backbones. """
233
234#         self.backbone_valid_task_ids: set[int] = set()
235#         r"""The list of task IDs that have valid backbones."""
236
237#     def on_train_start(self):
238#         r"""At the start of training for current task `self.task_id`, load the original backbone state dict to create a new independent backbone for the current task."""
239#         self.backbone.load_state_dict(self.original_backbone_state_dict)
240
241#     def on_train_end(self) -> None:
242#         r"""Save the trained independent backbone for `self.task_id`."""
243#         self.backbones[self.task_id] = deepcopy(self.backbone)
244#         self.backbone_valid_task_ids.add(self.task_id)
245
246#     def test_step(
247#         self, batch: DataLoader, batch_idx: int, dataloader_idx: int = 0
248#     ) -> dict[str, Tensor]:
249#         r"""Test step for current task `self.task_id`, which tests for all seen tasks indexed by `dataloader_idx`.
250
251#         **Args:**
252#         - **batch** (`Any`): a batch of test data.
253#         - **dataloader_idx** (`int`): the task ID of seen tasks to be tested. A default value of 0 is given otherwise the LightningModule will raise a `RuntimeError`.
254
255#         **Returns:**
256#         - **outputs** (`dict[str, Tensor]`): a dictionary contains loss and other metrics from this test step. Keys (`str`) are the metrics names, and values (`Tensor`) are the metrics.
257#         """
258#         test_task_id = self.get_test_task_id_from_dataloader_idx(dataloader_idx)
259
260#         x, y = batch
261#         backbone = self.backbones[
262#             test_task_id
263#         ]  # use the corresponding independenet backbone for the test task
264#         feature, _ = backbone(x, stage="test", task_id=test_task_id)
265#         logits = self.heads(feature, test_task_id)
266#         # use the corresponding head to test (instead of the current task `self.task_id`)
267#         loss_cls = self.criterion(logits, y)
268#         acc = (logits.argmax(dim=1) == y).float().mean()
269
270#         # Return metrics for lightning loggers callback to handle at `on_test_batch_end()`
271#         return {
272#             "loss_cls": loss_cls,
273#             "acc": acc,
274#         }
275
276
277# class UnlearnableIndependent(UnlearnableCLAlgorithm, Independent):
278#     r"""Unlearnable Independent learning algorithm.
279
280#     This is a variant of Independent that supports unlearning. It has the same functionality as Independent, but it also supports unlearning requests and permanent tasks.
281#     """
282
283#     def __init__(
284#         self,
285#         backbone: CLBackbone,
286#         heads: HeadsTIL | HeadsCIL | HeadDIL,
287#         non_algorithmic_hparams: dict[str, Any] = {},
288#         disable_unlearning: bool = False,
289#     ) -> None:
290#         r"""Initialize the Independent algorithm with the network. It has no additional hyperparameters.
291
292#         **Args:**
293#         - **backbone** (`CLBackbone`): backbone network.
294#         - **heads** (`HeadsTIL` | `HeadsCIL` | `HeadDIL`): output heads.
295#         - **disable_unlearning** (`bool`): whether to disable unlearning. This is used in reference experiments following continual learning pipeline. Default is `False`.
296
297#         """
298#         super().__init__(
299#             backbone=backbone,
300#             heads=heads,
301#             non_algorithmic_hparams=non_algorithmic_hparams,
302#             disable_unlearning=disable_unlearning,
303#         )
304
305#     def aggregated_backbone_output(self, input: Tensor) -> Tensor:
306#         r"""Get the aggregated backbone output for the input data. All parts of backbones should be aggregated together.
307
308#         This output feature is used for measuring unlearning metrics, such as Distribution Distance (DD). An aggregated output involving every part of the backbone is needed to ensure the fairness of the metric.
309
310#         **Args:**
311#         - **input** (`Tensor`): the input tensor from data.
312
313#         **Returns:**
314#         - **output** (`Tensor`): the aggregated backbone output tensor.
315#         """
316#         feature = 0
317
318#         for t in self.backbone_valid_task_ids:
319#             feature_t = self.backbones[t](input, stage="unlearning_test")[0]
320#             feature += feature_t
321#         feature = feature / len(self.backbone_valid_task_ids)
322
323#         return feature