clarena.cl_algorithms.wsn

  1r"""
  2The submodule in `cl_algorithms` for [WSN (Winning Subnetworks)](https://proceedings.mlr.press/v162/kang22b/kang22b.pdf) algorithm.
  3"""
  4
  5__all__ = ["WSN"]
  6
  7import logging
  8from typing import Any
  9
 10import torch
 11from torch import Tensor
 12from torch.utils.data import DataLoader
 13
 14from clarena.backbones import WSNMaskBackbone
 15from clarena.cl_algorithms import CLAlgorithm
 16from clarena.heads import HeadDIL, HeadsTIL
 17
 18# always get logger for built-in logging in each module
 19pylogger = logging.getLogger(__name__)
 20
 21
 22class WSN(CLAlgorithm):
 23    r"""[WSN (Winning Subnetworks)](https://proceedings.mlr.press/v162/kang22b/kang22b.pdf) algorithm.
 24
 25    An architecture-based continual learning approach that trains learnable parameter-wise scores and selects the most scored c% of network parameters per task.
 26    """
 27
 28    def __init__(
 29        self,
 30        backbone: WSNMaskBackbone,
 31        heads: HeadsTIL | HeadDIL,
 32        mask_percentage: float,
 33        parameter_score_init_mode: str = "default",
 34        non_algorithmic_hparams: dict[str, Any] = {},
 35        **kwargs,
 36    ) -> None:
 37        r"""Initialize the WSN algorithm with the network.
 38
 39        **Args:**
 40        - **backbone** (`WSNMaskBackbone`): must be a backbone network with the WSN mask mechanism.
 41        - **heads** (`HeadsTIL` | `HeadDIL`): output heads. WSN supports TIL (Task-Incremental Learning) and DIL (Domain-Incremental Learning).
 42        - **mask_percentage** (`float`): the percentage $c\%$ of parameters to be used for each task. See Sec. 3 and Eq. (4) in the [WSN paper](https://proceedings.mlr.press/v162/kang22b/kang22b.pdf).
 43        - **parameter_score_init_mode** (`str`): the initialization mode for parameter scores, must be one of:
 44            1. 'default': the default initialization in the original WSN code.
 45            2. 'N01': standard normal distribution $N(0, 1)$.
 46            3. 'U01': uniform distribution $U(0, 1)$.
 47        - **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.
 48        - **kwargs**: Reserved for multiple inheritance.
 49
 50        """
 51        super().__init__(
 52            backbone=backbone,
 53            heads=heads,
 54            non_algorithmic_hparams=non_algorithmic_hparams,
 55            **kwargs,
 56        )
 57
 58        self.mask_percentage: float = mask_percentage
 59        r"""The percentage of parameters to be used for each task."""
 60        self.parameter_score_init_mode: str = parameter_score_init_mode
 61        r"""The parameter score initialization mode."""
 62
 63        # save additional algorithmic hyperparameters
 64        self.save_hyperparameters(
 65            "mask_percentage",
 66            "parameter_score_init_mode",
 67        )
 68
 69        self.weight_masks: dict[int, dict[str, Tensor]] = {}
 70        r"""The binary weight mask of each previous task percentile-gated from the weight score. Keys are task IDs and values are the corresponding mask. Each mask is a dict where keys are layer names and values are the binary mask tensor for the layer. The mask tensor has the same size (output features, input features) as weight."""
 71        self.bias_masks: dict[int, dict[str, Tensor]] = {}
 72        r"""The binary bias mask of each previous task percentile-gated from the bias score. Keys are task IDs and values are the corresponding mask. Each mask is a dict where keys are layer names and values are the binary mask tensor for the layer. The mask tensor has the same size (output features, ) as bias. If the layer doesn't have bias, it is `None`."""
 73
 74        self.cumulative_weight_mask_for_previous_tasks: dict[str, Tensor] = {}
 75        r"""The cumulative binary weight mask $\mathbf{M}_{t-1}$ of previous tasks $1, \cdots, t-1$, percentile-gated from the weight score. It is a dict where keys are layer names and values are the binary mask tensors for the layers. The mask tensor has the same size (output features, input features) as weight."""
 76        self.cumulative_bias_mask_for_previous_tasks: dict[str, dict[str, Tensor]] = {}
 77        r"""The cumulative binary bias mask $\mathbf{M}_{t-1}$ of previous tasks $1, \cdots, t-1$, percentile-gated from the bias score. It is a dict where keys are layer names and values are the binary mask tensor for the layer. The mask tensor has the same size (output features, ) as bias. If the layer doesn't have bias, it is `None`."""
 78
 79        # set manual optimization
 80        self.automatic_optimization = False
 81
 82        WSN.sanity_check(self)
 83
 84    def sanity_check(self) -> None:
 85        r"""Sanity check."""
 86
 87        # check the backbone and heads
 88        if not isinstance(self.backbone, WSNMaskBackbone):
 89            raise ValueError("The backbone should be an instance of WSNMaskBackbone.")
 90        if not isinstance(self.heads, HeadsTIL):
 91            raise ValueError("The heads should be an instance of `HeadsTIL`.")
 92
 93        # check the mask percentage
 94        if not (0 < self.mask_percentage <= 1):
 95            raise ValueError(
 96                f"Mask percentage should be in (0, 1], but got {self.mask_percentage}."
 97            )
 98
 99    def on_train_start(self) -> None:
100        r"""Initialize the parameter scores before training the next task and initialize the cumulative masks at the beginning of the first task."""
101
102        self.backbone.initialize_parameter_score(mode=self.parameter_score_init_mode)
103
104        # initialize the cumulative mask at the beginning of the first task. This should not be called in `__init__()` because `self.device` is not available at that time.
105        if self.task_id == 1:
106            for layer_name in self.backbone.weighted_layer_names:
107                layer = self.backbone.get_layer_by_name(
108                    layer_name
109                )  # get the layer by its name
110
111                self.cumulative_weight_mask_for_previous_tasks[layer_name] = (
112                    torch.zeros_like(layer.weight).to(self.device)
113                )
114                if layer.bias is not None:
115                    self.cumulative_bias_mask_for_previous_tasks[layer_name] = (
116                        torch.zeros_like(layer.bias).to(self.device)
117                    )
118                else:
119                    self.cumulative_bias_mask_for_previous_tasks[layer_name] = None
120                # the cumulative mask $\mathrm{M}_{t-1}$ is initialized as zeros mask ($t = 1$)
121
122    def clip_grad_by_mask(
123        self,
124    ) -> None:
125        r"""Clip the gradients by the cumulative masks. The gradients are multiplied by (1 - cumulative_previous_mask) to keep previously masked parameters fixed. See Eq. (4) in the [WSN paper](https://proceedings.mlr.press/v162/kang22b/kang22b.pdf)."""
126
127        for layer_name in self.backbone.weighted_layer_names:
128            layer = self.backbone.get_layer_by_name(layer_name)
129
130            layer.weight.grad.data *= (
131                1 - self.cumulative_weight_mask_for_previous_tasks[layer_name]
132            )
133            if layer.bias is not None:
134                layer.bias.grad.data *= (
135                    1 - self.cumulative_bias_mask_for_previous_tasks[layer_name]
136                )
137
138    def forward(
139        self,
140        input: torch.Tensor,
141        stage: str,
142        task_id: int | None = None,
143    ) -> tuple[Tensor, dict[str, Tensor]]:
144        r"""The forward pass for data from task `task_id`. Note that it is nothing to do with `forward()` method in `nn.Module`.
145
146        **Args:**
147        - **input** (`Tensor`): the input tensor from data.
148        - **stage** (`str`): the stage of the forward pass, should be one of:
149            1. 'train': training stage.
150            2. 'validation': validation stage.
151            3. 'test': testing stage.
152        - **task_id** (`int` | `None`): the task ID where the data are from. If the stage is 'train' or 'validation', it should be the current task `self.task_id`. If the stage is 'test', it could be from any seen task (TIL uses the provided task IDs for testing).
153
154        **Returns:**
155        - **logits** (`Tensor`): the output logits tensor.
156        - **weight_mask** (`dict[str, Tensor]`): the weight mask for the current task. Key (`str`) is layer name, value (`Tensor`) is the mask tensor. The mask tensor has same (output features, input features) as weight.
157        - **bias_mask** (`dict[str, Tensor]`): the bias mask for the current task. Key (`str`) is layer name, value (`Tensor`) is the mask tensor. The mask tensor has same (output features, ) as bias. If the layer doesn't have bias, it is `None`.
158        - **activations** (`dict[str, Tensor]`): the hidden features (after activation) in each weighted layer. Key (`str`) is the weighted layer name, value (`Tensor`) is the hidden feature tensor. This is used for the continual learning algorithms that need to use the hidden features for various purposes.
159        """
160        feature, weight_mask, bias_mask, activations = self.backbone(
161            input,
162            stage=stage,
163            mask_percentage=self.mask_percentage,
164            test_mask=(
165                (self.weight_masks[task_id], self.bias_masks[task_id])
166                if stage == "test"
167                else None
168            ),
169        )
170        logits = self.heads(feature, task_id)
171
172        return (
173            logits
174            if self.if_forward_func_return_logits_only
175            else (logits, weight_mask, bias_mask, activations)
176        )
177
178    def training_step(self, batch: Any) -> dict[str, Tensor]:
179        r"""Training step for current task `self.task_id`.
180
181        **Args:**
182        - **batch** (`Any`): a batch of training data.
183
184        **Returns:**
185        - **outputs** (`dict[str, Tensor]`): a dictionary containing loss and other metrics from this training step. Keys (`str`) are the metrics names, and values (`Tensor`) are the metrics. Must include the key 'loss' which is total loss in the case of automatic optimization, according to PyTorch Lightning docs. For WSN, it includes 'weight_mask' and 'bias_mask' for logging.
186        """
187        x, y = batch
188
189        # zero the gradients before forward pass in manual optimization mode
190        opt = self.optimizers()
191        opt.zero_grad()
192
193        # classification loss
194        logits, weight_mask, bias_mask, activations = self.forward(
195            x, stage="train", task_id=self.task_id
196        )
197        loss_cls = self.criterion(logits, y)
198
199        # total loss
200        loss = loss_cls
201
202        # backward step (manually)
203        self.manual_backward(loss)  # calculate the gradients
204        # WSN hard-clips gradients using the cumulative masks. See Eq. (4) in the [WSN paper](https://proceedings.mlr.press/v162/kang22b/kang22b.pdf).
205        self.clip_grad_by_mask()
206
207        # update parameters with the modified gradients
208        opt.step()
209
210        # predicted labels
211        preds = logits.argmax(dim=1)
212
213        # accuracy of the batch
214        acc = (preds == y).float().mean()
215
216        return {
217            "preds": preds,
218            "loss": loss,  # return loss is essential for training step, or backpropagation will fail
219            "loss_cls": loss_cls,
220            "acc": acc,
221            "activations": activations,
222            "weight_mask": weight_mask,  # return other metrics for Lightning loggers callback to handle at `on_train_batch_end()`
223            "bias_mask": bias_mask,
224        }
225
226    def on_train_end(self) -> None:
227        r"""Store the weight and bias masks and update the cumulative masks after training the task."""
228
229        # get the weight and bias mask for the current task
230        weight_mask_t = {}
231        bias_mask_t = {}
232        for layer_name in self.backbone.weighted_layer_names:
233            layer = self.backbone.get_layer_by_name(layer_name)
234
235            weight_mask_t[layer_name] = self.backbone.gate_fn.apply(
236                self.backbone.weight_score_t[layer_name].weight, self.mask_percentage
237            )
238            if layer.bias is not None:
239                bias_mask_t[layer_name] = self.backbone.gate_fn.apply(
240                    self.backbone.bias_score_t[layer_name].weight.squeeze(
241                        0
242                    ),  # from (1, output_dim) to (output_dim, )
243                    self.mask_percentage,
244                )
245            else:
246                bias_mask_t[layer_name] = None
247
248        # store the weight and bias mask for the current task
249        self.weight_masks[self.task_id] = weight_mask_t
250        self.bias_masks[self.task_id] = bias_mask_t
251
252        # update the cumulative mask
253        for layer_name in self.backbone.weighted_layer_names:
254            layer = self.backbone.get_layer_by_name(layer_name)
255
256            self.cumulative_weight_mask_for_previous_tasks[layer_name] = torch.max(
257                self.cumulative_weight_mask_for_previous_tasks[layer_name],
258                weight_mask_t[layer_name],
259            )
260            if layer.bias is not None:
261                print(
262                    self.cumulative_bias_mask_for_previous_tasks[layer_name].shape,
263                    bias_mask_t[layer_name].shape,
264                )
265                self.cumulative_bias_mask_for_previous_tasks[layer_name] = torch.max(
266                    self.cumulative_bias_mask_for_previous_tasks[layer_name],
267                    bias_mask_t[layer_name],
268                )
269            else:
270                self.cumulative_bias_mask_for_previous_tasks[layer_name] = None
271
272        print(self.cumulative_bias_mask_for_previous_tasks)
273
274    def validation_step(self, batch: Any) -> dict[str, Tensor]:
275        r"""Validation step for current task `self.task_id`.
276
277        **Args:**
278        - **batch** (`Any`): a batch of validation data.
279
280        **Returns:**
281        - **outputs** (`dict[str, Tensor]`): a dictionary contains loss and other metrics from this validation step. Keys (`str`) are the metrics names, and values (`Tensor`) are the metrics.
282        """
283        x, y = batch
284        logits, _, _, _ = self.forward(x, stage="validation", task_id=self.task_id)
285        loss_cls = self.criterion(logits, y)
286        preds = logits.argmax(dim=1)
287        acc = (preds == y).float().mean()
288
289        return {
290            "loss_cls": loss_cls,
291            "acc": acc,  # return metrics for Lightning loggers callback to handle at `on_validation_batch_end()`
292            "preds": preds,
293        }
294
295    def test_step(
296        self, batch: DataLoader, batch_idx: int, dataloader_idx: int = 0
297    ) -> dict[str, Tensor]:
298        r"""Test step for current task `self.task_id`, which tests for all seen tasks indexed by `dataloader_idx`.
299
300        **Args:**
301        - **batch** (`Any`): a batch of test data.
302        - **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`.
303
304        **Returns:**
305        - **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.
306        """
307        test_task_id = self.get_test_task_id_from_dataloader_idx(dataloader_idx)
308
309        x, y = batch
310        logits, _, _, _ = self.forward(
311            x,
312            stage="test",
313            task_id=test_task_id,
314        )  # use the corresponding head and mask to test (instead of the current task `self.task_id`)
315        loss_cls = self.criterion(logits, y)
316        preds = logits.argmax(dim=1)
317        acc = (preds == y).float().mean()
318
319        return {
320            "loss_cls": loss_cls,
321            "acc": acc,  # return metrics for Lightning loggers callback to handle at `on_test_batch_end()`
322            "preds": preds,
323        }