diff --git a/src/capymoa/ocl/strategy/__init__.py b/src/capymoa/ocl/strategy/__init__.py index 8ab23c16..8820d526 100644 --- a/src/capymoa/ocl/strategy/__init__.py +++ b/src/capymoa/ocl/strategy/__init__.py @@ -8,5 +8,6 @@ from . import l2p from ._ewc import EWC from ._si import SI +from ._lwf import LWF -__all__ = ["ExperienceReplay", "SLDA", "NCM", "GDumb", "RAR", "l2p", "EWC", "SI"] +__all__ = ["ExperienceReplay", "SLDA", "NCM", "GDumb", "RAR", "l2p", "EWC", "SI", "LWF"] diff --git a/src/capymoa/ocl/strategy/_lwf.py b/src/capymoa/ocl/strategy/_lwf.py new file mode 100644 index 00000000..d5df3b42 --- /dev/null +++ b/src/capymoa/ocl/strategy/_lwf.py @@ -0,0 +1,118 @@ +from copy import deepcopy +from typing import Optional + +import torch +from torch import Tensor, nn + +from capymoa.base import BatchClassifier +from capymoa.ocl.events import Dispatcher, Handler +from capymoa.ocl.evaluation.events import TrainTaskBegin +from capymoa.ocl.util._optim import reset_optimizer_state +from capymoa.ocl.util.functional import hinton_distillation_loss +from capymoa.stream._stream import Schema + + +class LWF(BatchClassifier, nn.Module, Handler): + """Learning Without Forgetting (LwF). + + LwF [#f1]_ is a regularisation-based continual learning strategy that distils + predictions from a frozen teacher snapshot of the previous task while learning the + current task. + + .. [#f1] Li, Z., & Hoiem, D. (2016). Learning without forgetting. CoRR, + abs/1606.09282. http://arxiv.org/abs/1606.09282 + """ + + def __init__( + self, + schema: Schema, + model: torch.nn.Module, + optimiser: torch.optim.Optimizer, + alpha: float = 1.0, + temperature: float = 2.0, + device: torch.device = torch.device("cpu"), + ) -> None: + """Construct an LWF learner. + + :param schema: Stream schema used by the classifier interface. + :param model: Torch model that outputs class logits. + :param optimiser: Optimiser used to update ``model`` parameters. + :param alpha: Weight of the distillation loss term. + :param temperature: Distillation temperature. + :param device: Compute device. + :raises ValueError: If ``alpha`` is negative or ``temperature`` is not + positive. + """ + super().__init__(schema, 0) + nn.Module.__init__(self) + if alpha < 0: + raise ValueError("alpha must be non-negative.") + if temperature <= 0: + raise ValueError("temperature must be greater than zero.") + + self.device = device + + self._alpha = alpha + self._temperature = temperature + + self._optimiser = optimiser + self._model = model + self._criterion = torch.nn.CrossEntropyLoss() + + self._teacher: Optional[torch.nn.Module] = None + self._train_task = 0 + + # Move all model parameters and buffers to the specified device + self.to(device) + + def batch_train(self, x: Tensor, y: Tensor) -> None: + self._model.train() + self._optimiser.zero_grad() + + student_logits = self._model(x) + task_loss = self._criterion(student_logits, y) + total_loss = task_loss + self._alpha * self._distillation_loss( + x, student_logits + ) + + total_loss.backward() + self._optimiser.step() + + @torch.no_grad() + def batch_predict_proba(self, x: Tensor) -> Tensor: + self._model.eval() + y_hat = self._model(x) + return torch.softmax(y_hat, dim=1) + + def attach_with(self, source: Dispatcher) -> "LWF": + source.subscribe(TrainTaskBegin, self._on_train_task_begin) + return self + + def _on_train_task_begin(self, event: TrainTaskBegin) -> None: + reset_optimizer_state(self._optimiser) + if event.train_task > 0: + self._teacher = ( + deepcopy(self._model).to(self.device).eval().requires_grad_(False) + ) + self._train_task = event.train_task + + @torch.no_grad() + def _teacher_forward(self, x: Tensor) -> Tensor: + if self._teacher is None: + raise RuntimeError("Teacher model is not available before task 1.") + return self._teacher(x) + + def _distillation_loss(self, x: Tensor, student_logits: Tensor) -> Tensor: + if self._teacher is None: + return torch.tensor(0.0, device=self.device) + + teacher_logits = self._teacher_forward(x) + + return hinton_distillation_loss( + teacher_logits=teacher_logits, + student_logits=student_logits, + temperature=self._temperature, + ) + + def __str__(self) -> str: + return f"LWF(alpha={self._alpha}, temperature={self._temperature})" diff --git a/src/capymoa/ocl/util/functional.py b/src/capymoa/ocl/util/functional.py new file mode 100644 index 00000000..c10239e7 --- /dev/null +++ b/src/capymoa/ocl/util/functional.py @@ -0,0 +1,35 @@ +"""A collection of functional utilities for OCL.""" + +from torch import Tensor +from torch.nn.functional import kl_div, log_softmax + + +def hinton_distillation_loss( + teacher_logits: Tensor, student_logits: Tensor, temperature: float = 1.0 +) -> Tensor: + """Hinton's distillation loss [#f1]_ . + + .. math:: + L_{KD} = T^2 KL(softmax(z_s / T), softmax(z_t / T)) + + where :math:`T` is the temperature, :math:`z_s` are the student logits, and + :math:`z_t` are the teacher logits. + + .. [#f1] Hinton, G., Vinyals, O., & Dean, J. (2015). Distilling the Knowledge in a + Neural Network. arXiv:1503.02531 [Cs, Stat]. http://arxiv.org/abs/1503.02531 + + :param teacher_logits: Teacher logits of shape ``(batch_size, num_classes)``. + :param student_logits: Student logits of shape ``(batch_size, num_classes)``. + :param temperature: Temperature for distillation. Higher values produce softer + probability distributions. + :return: The distillation loss as a scalar tensor. + """ + return ( + kl_div( + log_softmax(student_logits / temperature, dim=1), # Soft predictions + log_softmax(teacher_logits / temperature, dim=1), # Soft targets + log_target=True, + reduction="batchmean", # Mathematically correct unlike the default + ) + * temperature**2 + ) diff --git a/tests/ocl/test_strategy.py b/tests/ocl/test_strategy.py index 83cf1317..735f0535 100644 --- a/tests/ocl/test_strategy.py +++ b/tests/ocl/test_strategy.py @@ -14,7 +14,7 @@ from capymoa.ocl.datasets import TinySplitMNIST # noqa: E402 from capymoa.ocl.evaluation import ocl_train_eval_loop # noqa: E402 -from capymoa.ocl.strategy import ExperienceReplay, SLDA, NCM, GDumb, RAR, EWC, SI # noqa: E402 +from capymoa.ocl.strategy import ExperienceReplay, SLDA, NCM, GDumb, RAR, EWC, SI, LWF # noqa: E402 from capymoa.stream import Schema # noqa: E402 import torch # noqa: E402 @@ -156,6 +156,11 @@ def _new_rar(schema): Result(96.5, 80.8, 34.4), task_mask=True, ), + Case( + "LWF", + new_constructor(LWF, lr=0.10, alpha=4.66, temperature=1.67), + Result(36.49, 25.09, 17.59), + ), ]