22class CLAlgorithm(LightningModule):
 23    """
 24    The base class of continual learning algorithms, inherited from `LightningModule`.
 25    """
 26
 27    def __init__(
 28        self,
 29        backbone: CLBackbone,
 30        heads: HeadsTIL | HeadsCIL,
 31    ) -> None:
 32        """Initialise the CL algorithm with the network.
 33
 34        Args:
 35        - **backbone** (`CLBackbone`): backbone network.
 36        - **heads** (`HeadsTIL` | `HeadsCIL`): output heads.
 37        """
 38        super().__init__()
 39
 40        self.backbone: CLBackbone = backbone
 41        """Store the backbone network."""
 42        self.heads: HeadsTIL | HeadsCIL = heads
 43        """Store the output heads."""
 44        self.optimizer: Optimizer
 45        """Store the optimizer object (partially initialised) for the backpropagation of task `self.task_id`. Will be equipped with parameters in `configure_optimizers()`."""
 46        self.criterion = nn.CrossEntropyLoss()
 47        """The loss function bewteen the output logits and the target labels. Default is cross-entropy loss."""
 48
 49        self.task_id: int
 50        """Task ID counter indicating which task is being processed. Self updated during the task loop."""
 51
 52        self.sanity_check()
 53
 54    def sanity_check(self) -> None:
 55        """Check the sanity of the arguments.
 56
 57        **Raises:**
 58        - **ValueError**: if the `output_dim` of backbone network is not equal to the `input_dim` of CL heads.
 59        """
 60        if self.backbone.output_dim != self.heads.input_dim:
 61            raise ValueError(
 62                "The output_dim of backbone network should be equal to the input_dim of CL heads!"
 63            )
 64
 65    def setup_task_id(
 66        self, task_id: int, num_classes_t: int, optimizer: Optimizer
 67    ) -> None:
 68        """Set up which task's dataset the CL experiment is on. This must be done before `forward()` method is called.
 69
 70        **Args:**
 71        - **task_id** (`int`): the target task ID.
 72        - **num_classes_t** (`int`): the number of classes in the task.
 73        - **optimizer** (`Optimizer`): the optimizer object (partially initialised) for the task `self.task_id`.
 74        """
 75        self.task_id = task_id
 76        self.heads.setup_task_id(task_id, num_classes_t)
 77        self.optimizer = optimizer
 78
 79    @override
 80    def forward(self, input: Tensor, task_id: int):
 81        """The default forward pass for data from task `task_id`. Note that it is nothing to do with `forward()` method in `nn.Module`.
 82
 83        **Args:**
 84        - **input** (`Tensor`): The input tensor from data.
 85        - **task_id** (`int`): the task ID where the data are from.
 86
 87        Returns:
 88        - The output logits tensor.
 89        """
 90        feature = self.backbone(input, task_id)
 91        logits = self.heads(feature, task_id)
 92        return logits
 93
 94    def configure_optimizers(self) -> Optimizer:
 95        """
 96        Configure optimizer hooks by Lightning.
 97        See [Lightning docs](https://lightning.ai/docs/pytorch/stable/common/lightning_module.html#configure-optimizers) for more details.
 98        """
 99        # finish partially initialised optimizer by specifying model parameters. The `parameters()` method of this `CLAlrogithm` (inherited from `LightningModule`) returns both backbone and heads parameters
100        return self.optimizer(params=self.parameters())