clarena.cl_algorithms.finetuning

 1"""
 2The submodule in `cl_algorithms` for Finetuning algorithm.
 3"""
 4
 5__all__ = ["Finetuning"]
 6
 7import logging
 8from typing import Any
 9
10import torch
11from torch.utils.data import DataLoader
12
13from clarena.cl_algorithms import CLAlgorithm
14
15# always get logger for built-in logging in each module
16pylogger = logging.getLogger(__name__)
17
18
19class Finetuning(CLAlgorithm):
20    """Finetuning algorithm.
21
22    It is the most naive way for task-incremental learning. It simply initialises the backbone from the last task when training new task.
23    """
24
25    def __init__(
26        self,
27        backbone: torch.nn.Module,
28        heads: torch.nn.Module,
29    ) -> None:
30        """Initialise the Finetuning algorithm with the network. It has no additional hyperparamaters.
31
32        **Args:**
33        - **backbone** (`CLBackbone`): backbone network.
34        - **heads** (`HeadsTIL` | `HeadsCIL`): output heads.
35        """
36        super().__init__(backbone=backbone, heads=heads)
37
38    def training_step(self, batch: Any):
39        """Training step for current task `self.task_id`.
40
41        **Args:**
42        - **batch** (`Any`): a batch of training data.
43        """
44        x, y = batch
45        logits = self.forward(x, self.task_id)
46        loss_cls = self.criterion(logits, y)
47        loss = loss_cls
48        acc = (logits.argmax(dim=1) == y).float().mean()
49
50        # Return loss is essential for training step, or backpropagation will fail
51        # Return other metrics for lightning loggers callback to handle at `on_train_batch_end()`
52        return {
53            "loss": loss,
54            "loss_cls": loss_cls,
55            "acc": acc,
56        }
57
58    def validation_step(self, batch: Any):
59        """Validation step for current task `self.task_id`.
60
61        **Args:**
62        - **batch** (`Any`): a batch of validation data.
63        """
64        x, y = batch
65        logits = self.forward(x, self.task_id)
66        loss_cls = self.criterion(logits, y)
67        acc = (logits.argmax(dim=1) == y).float().mean()
68
69        # Return metrics for lightning loggers callback to handle at `on_validation_batch_end()`
70        return {
71            "loss_cls": loss_cls,
72            "acc": acc,
73        }
74
75    def test_step(self, batch: DataLoader, batch_idx: int, dataloader_idx: int = 0):
76        """Test step for current task `self.task_id`, which tests for all seen tasks indexed by `dataloader_idx`.
77
78        **Args:**
79        - **batch** (`Any`): a batch of test data.
80        - **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`.
81        """
82        test_task_id = dataloader_idx + 1
83
84        x, y = batch
85        logits = self.forward(
86            x, test_task_id
87        )  # use the corresponding head to test (instead of the current task `self.task_id`)
88        loss_cls = self.criterion(logits, y)
89        acc = (logits.argmax(dim=1) == y).float().mean()
90
91        # Return metrics for lightning loggers callback to handle at `on_test_batch_end()`
92        return {
93            "loss_cls": loss_cls,
94            "acc": acc,
95        }