25189796363

Committed 30 Apr 2026 09:17PM UTC coverage: 75.524% (+0.6%) from 74.909%

Build # 25189796363

Build Type

Pull #355

github

Committed by

web-flow

Commit Message

Merge e4643751b into 2ef7fe73e

Pull Request Pull Request #355: feat(ocl): add LWF and DER

Coverage Stats

366 of 372 new or added lines in 19 files covered. (98.39%)

6 existing lines in 4 files now uncovered.

7452 of 9867 relevant lines covered (75.52%)

0.76 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

96.23

/src/capymoa/ocl/strategy/_ewc.py

from typing import Iterable, Iterator, Optional, Sequence, Tuple, Callable
from capymoa.stream._stream import Schema
from torch import Tensor, nn
import torch
from capymoa.base import BatchClassifier
from capymoa.ocl.base import TrainTaskAware, TestTaskAware
from capymoa.ocl.util._buffer import BufferList
from capymoa.ocl.util._replay import SlidingWindow
from torch.utils.data import DataLoader


def weighted_l2_reg(
    params: Iterable[Tensor],
    anchor_params: Iterable[Tensor],
    fisher_diagonals: Iterable[Tensor],
    device: torch.device,
) -> Tensor:
    """Compute an EWC-style weighted L2 regularisation term.

    :param params: Current model parameters.
    :param anchor_params: Reference parameters from a previous task.
    :param fisher_diagonals: Diagonal Fisher information weights.
    :param device: Device used for the accumulator tensor.
    :return: Weighted L2 penalty scaled by ``1/2``.
    """
    l2 = torch.tensor(0.0, device=device)
    for param, anchor_param, fisher_diag in zip(
        params, anchor_params, fisher_diagonals, strict=True
    ):
        assert param.shape == anchor_param.shape
        l2 += (fisher_diag * (param - anchor_param) ** 2).sum()
    return l2 / 2.0


def fd_init(model: torch.nn.Module) -> Sequence[Tensor]:
    """Initialise zero-valued Fisher diagonal tensors for a model.

    :param model: Model whose parameters define the Fisher diagonal shapes.
    :return: Zero tensors matching all model parameters.
    """
    return [torch.zeros_like(param) for param in model.parameters()]


def fd_accumulate(
    fisher_diagonals: Sequence[Tensor],
    parameters: Iterator[Tensor],
    alpha: Optional[float] = None,
) -> Sequence[Tensor]:
    """Accumulates the squared gradients into the Fisher diagonal estimates.

    :param fisher_diagonals: A sequence of tensors representing the current estimates of
        the Fisher diagonals.
    :param parameters: A sequence of model parameters whose gradients have been
        computed.
    :param alpha: Decay factor for the accumulated Fisher diagonals. A value of 1.0
        corresponds to standard EWC accumulation, while values less than 1.0 implement
        a decay as in Online EWC.
    :return: Updated sequence of tensors representing the accumulated Fisher diagonals.
    """
    for fisher_diag, param in zip(fisher_diagonals, parameters, strict=True):
        if param.grad is None:
            raise ValueError(
                "Parameter gradients must be computed before updating Fisher diagonals."
            )
        if alpha is not None:
            fisher_diag.mul_(alpha).add_(param.grad.data.pow(2), alpha=(1 - alpha))
        else:
            fisher_diag.add_(param.grad.data.pow(2))
    return fisher_diagonals


def fd_compute(
    model: torch.nn.Module,
    forward_fn: Callable[[Tensor], Tensor],
    dataloader: DataLoader[Tuple[Tensor, Tensor]],
    device: torch.device,
    criterion: torch.nn.Module,
) -> Sequence[Tensor]:
    """Compute module fisher diagonals.

    :param model: A PyTorch classifier model.
    :param dataloader: A PyTorch dataloader for a classification task, yielding batches
        of (inputs, labels).
    :param device: Compute device.
    :param criterion: The loss function to use.
    :return: A sequence of tensors representing the computed Fisher diagonals.
    """
    model = model.eval().to(device)
    criterion = criterion.eval().to(device)

    fisher_diagonals = fd_init(model)
    for inputs, labels in dataloader:
        model.zero_grad()
        inputs, labels = inputs.to(device), labels.to(device)
        outputs = forward_fn(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        fisher_diagonals = fd_accumulate(fisher_diagonals, model.parameters())
    # Average the accumulated squared gradients over the number of samples
    fisher_diagonals = [
        fisher_diag / len(dataloader) for fisher_diag in fisher_diagonals
    ]
    return fisher_diagonals


class EWC(BatchClassifier, nn.Module, TrainTaskAware, TestTaskAware):
    """Elastic Weight Consolidation learner.

    Elastic Weight Consolidation (EWC) is a regularisation-based continual learning
    strategy that mitigates catastrophic forgetting by penalising changes to important
    parameters for previous tasks [#f1]_. We incorporate Online EWC-style [#f2]_ updates
    to the Fisher diagonals, which decay the importance of previous tasks' parameters
    over time based on the ``gamma`` hyperparameter.

    Usually the EWC strategy has access to the entire active task's data when estimating
    the Fisher diagonals, but instead we use a replay buffer to approximate the active
    task distribution.

    ..  [#f1] Kirkpatrick, J., Pascanu, R., Rabinowitz, N., Veness, J., Desjardins, G.,
        Rusu, A. A., Milan, K., Quan, J., Ramalho, T., Grabska-Barwinska, A., Hassabis,
        D., Clopath, C., Kumaran, D., & Hadsell, R. (2017). Overcoming catastrophic
        forgetting in neural networks. Proceedings of the National Academy of Sciences,
        114(13), 3521–3526. https://doi.org/10.1073/pnas.1611835114

    ..  [#f2] Schwarz, J., Czarnecki, W., Luketina, J., Grabska-Barwinska, A., Teh, Y.
        W., Pascanu, R., & Hadsell, R. (2018). Progress & Compress: A scalable framework
        for continual learning. In J. G. Dy & A. Krause (Eds.), Proceedings of the 35th
        International Conference on Machine Learning, ICML 2018, Stockholmsmässan,
        Stockholm, Sweden, July 10-15, 2018 (Vol. 80, pp. 4535–4544). PMLR.
        http://proceedings.mlr.press/v80/schwarz18a.html
    """

    def __init__(
        self,
        schema: Schema,
        model: torch.nn.Module,
        optimiser: torch.optim.Optimizer,
        lambda_: float,
        fim_buffer: int = 256,
        fim_batch_size: int = 32,
        device: torch.device = torch.device("cpu"),
        mask_test: bool = False,
        mask_train: bool = False,
        gamma: float = 1.0,
        task_mask: Optional[Tensor] = None,
    ) -> None:
        """Construct an EWC 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 lambda_: Weight of the EWC regularisation term.
        :param fim_buffer: Replay window size for Fisher estimation.
        :param fim_batch_size: Mini-batch size used when estimating Fisher diagonals.
        :param device: Compute device.
        :param mask_test: Whether to apply per-task masking during testing. This is a
            task incremental scenario.
        :param mask_train: Whether to apply per-task masking during training. This is
            also known as the labels trick.
        :param task_mask: Optional per-task mask applied to output logits.
        :raises ValueError: If task-specific masking is requested without ``task_mask``.
        """
        super().__init__(schema, 0)
        nn.Module.__init__(self)
        if (mask_train or mask_test) and task_mask is None:
            raise ValueError(
                "Task schedule must be provided for task incremental or labels trick scenarios."
            )
        self.device = device

        # Hyperparameters
        self._lambda = lambda_
        self._gamma = gamma
        self._fd_batch_size = fim_batch_size
        self._mask_train = mask_train
        self._mask_test = mask_test

        # Modules
        self._optimiser = optimiser
        self._model = model
        self._criterion = torch.nn.CrossEntropyLoss()
        self._buffer = SlidingWindow.new_xybuffer(fim_buffer, schema.shape)

        # Buffers for anchoring the model
        self._anchor_params = BufferList(
            [param.clone().detach() for param in model.parameters()]
        )
        self._fisher_diags = BufferList(
            [torch.zeros_like(param) for param in model.parameters()]
        )

        # Task tracking
        self._train_task = 0
        self._test_task = 0
        if task_mask is None:
            self._task_mask = None
        else:
            self._task_mask = nn.Buffer(task_mask)

        # Move all model parameters and buffers to the specified device
        self.to(device)

    def batch_train(self, x: Tensor, y: Tensor) -> None:
        self._buffer.update(x=x, y=y)
        self._model.train()
        self._optimiser.zero_grad()
        y_hat = self._train_forward(x)
        loss = self._criterion(y_hat, y)
        total_loss = loss + self._lambda * self._regularisation_loss()
        total_loss.backward()
        self._optimiser.step()

    @torch.no_grad()
    def batch_predict_proba(self, x: Tensor) -> Tensor:
        self._model.eval()
        y_hat = self._test_forward(x)
        return torch.softmax(y_hat, dim=1)

    def on_train_task(self, task_id: int) -> None:
        if task_id > 0:
            self._update_fisher_diags()
            self._update_anchor_params()
        self._train_task = task_id

    def on_test_task(self, task_id: int) -> None:
        self._test_task = task_id

    def _update_fisher_diags(self) -> None:
        """Estimate and accumulate Fisher diagonals from the replay buffer."""
        dataset = self._buffer.dataset_view()
        dataloader = DataLoader(dataset, batch_size=self._fd_batch_size, shuffle=False)
        task_fisher_diags = fd_compute(
            self._model,
            self._train_forward,
            dataloader,  # type: ignore
            self.device,
            self._criterion,
        )
        # Update the fisher diagonals buffer with the computed values
        for i in range(len(self._fisher_diags)):
            self._fisher_diags[i].mul_(self._gamma).add_(task_fisher_diags[i])

    def _update_anchor_params(self) -> None:
        """Update anchored parameters to the current model weights."""
        for param, anchor_param in zip(
            self._model.parameters(), self._anchor_params, strict=True
        ):
            anchor_param.copy_(param.detach())

    def _test_forward(self, x: Tensor) -> Tensor:
        """Compute logits for inference, optionally applying a test-task mask."""
        y_hat = self._model(x)
        if self._task_mask is not None and self._mask_test:
            y_hat = self._task_mask[self._test_task] * y_hat
        return y_hat

    def _train_forward(self, x: Tensor) -> Tensor:
        """Compute logits for training, optionally applying a train-task mask."""
        y_hat = self._model(x)
        if self._task_mask is not None and self._mask_train:
            y_hat = self._task_mask[self._train_task] * y_hat
        return y_hat

    def _regularisation_loss(self) -> torch.Tensor:
        """Return the EWC regularisation loss for the current task."""
        if self._train_task < 1:
            return torch.tensor(0.0, device=self.device)
        return weighted_l2_reg(
            self._model.parameters(),
            self._anchor_params,
            self._fisher_diags,
            device=self.device,
        )

1	from typing import Iterable, Iterator, Optional, Sequence, Tuple, Callable	1✔
2	from capymoa.stream._stream import Schema	1✔
3	from torch import Tensor, nn	1✔
4	import torch	1✔
5	from capymoa.base import BatchClassifier	1✔
6	from capymoa.ocl.base import TrainTaskAware, TestTaskAware	1✔
7	from capymoa.ocl.util._buffer import BufferList	1✔
8	from capymoa.ocl.util._replay import SlidingWindow	1✔
9	from torch.utils.data import DataLoader	1✔
10
11
12	def weighted_l2_reg(	1✔
13	params: Iterable[Tensor],
14	anchor_params: Iterable[Tensor],
15	fisher_diagonals: Iterable[Tensor],
16	device: torch.device,
17	) -> Tensor:
18	"""Compute an EWC-style weighted L2 regularisation term.
19
20	:param params: Current model parameters.
21	:param anchor_params: Reference parameters from a previous task.
22	:param fisher_diagonals: Diagonal Fisher information weights.
23	:param device: Device used for the accumulator tensor.
24	:return: Weighted L2 penalty scaled by ``1/2``.
25	"""
26	l2 = torch.tensor(0.0, device=device)	1✔
27	for param, anchor_param, fisher_diag in zip(	1✔
28	params, anchor_params, fisher_diagonals, strict=True
29	):
30	assert param.shape == anchor_param.shape	1✔
31	l2 += (fisher_diag * (param - anchor_param) ** 2).sum()	1✔
32	return l2 / 2.0	1✔
33
34
35	def fd_init(model: torch.nn.Module) -> Sequence[Tensor]:	1✔
36	"""Initialise zero-valued Fisher diagonal tensors for a model.
37
38	:param model: Model whose parameters define the Fisher diagonal shapes.
39	:return: Zero tensors matching all model parameters.
40	"""
41	return [torch.zeros_like(param) for param in model.parameters()]	1✔
42
43
44	def fd_accumulate(	1✔
45	fisher_diagonals: Sequence[Tensor],
46	parameters: Iterator[Tensor],
47	alpha: Optional[float] = None,
48	) -> Sequence[Tensor]:
49	"""Accumulates the squared gradients into the Fisher diagonal estimates.
50
51	:param fisher_diagonals: A sequence of tensors representing the current estimates of
52	the Fisher diagonals.
53	:param parameters: A sequence of model parameters whose gradients have been
54	computed.
55	:param alpha: Decay factor for the accumulated Fisher diagonals. A value of 1.0
56	corresponds to standard EWC accumulation, while values less than 1.0 implement
57	a decay as in Online EWC.
58	:return: Updated sequence of tensors representing the accumulated Fisher diagonals.
59	"""
60	for fisher_diag, param in zip(fisher_diagonals, parameters, strict=True):	1✔
61	if param.grad is None:	1✔
62	raise ValueError(	×
63	"Parameter gradients must be computed before updating Fisher diagonals."
64	)
65	if alpha is not None:	1✔
66	fisher_diag.mul_(alpha).add_(param.grad.data.pow(2), alpha=(1 - alpha))	×
67	else:
68	fisher_diag.add_(param.grad.data.pow(2))	1✔
69	return fisher_diagonals	1✔
70
71
72	def fd_compute(	1✔
73	model: torch.nn.Module,
74	forward_fn: Callable[[Tensor], Tensor],
75	dataloader: DataLoader[Tuple[Tensor, Tensor]],
76	device: torch.device,
77	criterion: torch.nn.Module,
78	) -> Sequence[Tensor]:
79	"""Compute module fisher diagonals.
80
81	:param model: A PyTorch classifier model.
82	:param dataloader: A PyTorch dataloader for a classification task, yielding batches
83	of (inputs, labels).
84	:param device: Compute device.
85	:param criterion: The loss function to use.
86	:return: A sequence of tensors representing the computed Fisher diagonals.
87	"""
88	model = model.eval().to(device)	1✔
89	criterion = criterion.eval().to(device)	1✔
90
91	fisher_diagonals = fd_init(model)	1✔
92	for inputs, labels in dataloader:	1✔
93	model.zero_grad()	1✔
94	inputs, labels = inputs.to(device), labels.to(device)	1✔
95	outputs = forward_fn(inputs)	1✔
96	loss = criterion(outputs, labels)	1✔
97	loss.backward()	1✔
98	fisher_diagonals = fd_accumulate(fisher_diagonals, model.parameters())	1✔
99	# Average the accumulated squared gradients over the number of samples
100	fisher_diagonals = [	1✔
101	fisher_diag / len(dataloader) for fisher_diag in fisher_diagonals
102	]
103	return fisher_diagonals	1✔
104
105
106	class EWC(BatchClassifier, nn.Module, TrainTaskAware, TestTaskAware):	1✔
107	"""Elastic Weight Consolidation learner.
108
109	Elastic Weight Consolidation (EWC) is a regularisation-based continual learning
110	strategy that mitigates catastrophic forgetting by penalising changes to important
111	parameters for previous tasks [#f1]_. We incorporate Online EWC-style [#f2]_ updates
112	to the Fisher diagonals, which decay the importance of previous tasks' parameters
113	over time based on the ``gamma`` hyperparameter.
114
115	Usually the EWC strategy has access to the entire active task's data when estimating
116	the Fisher diagonals, but instead we use a replay buffer to approximate the active
117	task distribution.
118
119	.. [#f1] Kirkpatrick, J., Pascanu, R., Rabinowitz, N., Veness, J., Desjardins, G.,
120	Rusu, A. A., Milan, K., Quan, J., Ramalho, T., Grabska-Barwinska, A., Hassabis,
121	D., Clopath, C., Kumaran, D., & Hadsell, R. (2017). Overcoming catastrophic
122	forgetting in neural networks. Proceedings of the National Academy of Sciences,
123	114(13), 3521–3526. https://doi.org/10.1073/pnas.1611835114
124
125	.. [#f2] Schwarz, J., Czarnecki, W., Luketina, J., Grabska-Barwinska, A., Teh, Y.
126	W., Pascanu, R., & Hadsell, R. (2018). Progress & Compress: A scalable framework
127	for continual learning. In J. G. Dy & A. Krause (Eds.), Proceedings of the 35th
128	International Conference on Machine Learning, ICML 2018, Stockholmsmässan,
129	Stockholm, Sweden, July 10-15, 2018 (Vol. 80, pp. 4535–4544). PMLR.
130	http://proceedings.mlr.press/v80/schwarz18a.html
131	"""
132
133	def __init__(	1✔
134	self,
135	schema: Schema,
136	model: torch.nn.Module,
137	optimiser: torch.optim.Optimizer,
138	lambda_: float,
139	fim_buffer: int = 256,
140	fim_batch_size: int = 32,
141	device: torch.device = torch.device("cpu"),
142	mask_test: bool = False,
143	mask_train: bool = False,
144	gamma: float = 1.0,
145	task_mask: Optional[Tensor] = None,
146	) -> None:
147	"""Construct an EWC learner.
148
149	:param schema: Stream schema used by the classifier interface.
150	:param model: Torch model that outputs class logits.
151	:param optimiser: Optimiser used to update ``model`` parameters.
152	:param lambda_: Weight of the EWC regularisation term.
153	:param fim_buffer: Replay window size for Fisher estimation.
154	:param fim_batch_size: Mini-batch size used when estimating Fisher diagonals.
155	:param device: Compute device.
156	:param mask_test: Whether to apply per-task masking during testing. This is a
157	task incremental scenario.
158	:param mask_train: Whether to apply per-task masking during training. This is
159	also known as the labels trick.
160	:param task_mask: Optional per-task mask applied to output logits.
161	:raises ValueError: If task-specific masking is requested without ``task_mask``.
162	"""
163	super().__init__(schema, 0)	1✔
164	nn.Module.__init__(self)	1✔
165	if (mask_train or mask_test) and task_mask is None:	1✔
UNCOV 166	raise ValueError(	×
167	"Task schedule must be provided for task incremental or labels trick scenarios."
168	)
169	self.device = device	1✔
170
171	# Hyperparameters
172	self._lambda = lambda_	1✔
173	self._gamma = gamma	1✔
174	self._fd_batch_size = fim_batch_size	1✔
175	self._mask_train = mask_train	1✔
176	self._mask_test = mask_test	1✔
177
178	# Modules
179	self._optimiser = optimiser	1✔
180	self._model = model	1✔
181	self._criterion = torch.nn.CrossEntropyLoss()	1✔
182	self._buffer = SlidingWindow.new_xybuffer(fim_buffer, schema.shape)	1✔
183
184	# Buffers for anchoring the model
185	self._anchor_params = BufferList(	1✔
186	[param.clone().detach() for param in model.parameters()]
187	)
188	self._fisher_diags = BufferList(	1✔
189	[torch.zeros_like(param) for param in model.parameters()]
190	)
191
192	# Task tracking
193	self._train_task = 0	1✔
194	self._test_task = 0	1✔
195	if task_mask is None:	1✔
196	self._task_mask = None	1✔
197	else:
198	self._task_mask = nn.Buffer(task_mask)	1✔
199
200	# Move all model parameters and buffers to the specified device
201	self.to(device)	1✔
202
203	def batch_train(self, x: Tensor, y: Tensor) -> None:	1✔
204	self._buffer.update(x=x, y=y)	1✔
205	self._model.train()	1✔
206	self._optimiser.zero_grad()	1✔
207	y_hat = self._train_forward(x)	1✔
208	loss = self._criterion(y_hat, y)	1✔
209	total_loss = loss + self._lambda * self._regularisation_loss()	1✔
210	total_loss.backward()	1✔
211	self._optimiser.step()	1✔
212
213	@torch.no_grad()	1✔
214	def batch_predict_proba(self, x: Tensor) -> Tensor:	1✔
215	self._model.eval()	1✔
216	y_hat = self._test_forward(x)	1✔
217	return torch.softmax(y_hat, dim=1)	1✔
218
219	def on_train_task(self, task_id: int) -> None:	1✔
220	if task_id > 0:	1✔
221	self._update_fisher_diags()	1✔
222	self._update_anchor_params()	1✔
223	self._train_task = task_id	1✔
224
225	def on_test_task(self, task_id: int) -> None:	1✔
226	self._test_task = task_id	1✔
227
228	def _update_fisher_diags(self) -> None:	1✔
229	"""Estimate and accumulate Fisher diagonals from the replay buffer."""
230	dataset = self._buffer.dataset_view()	1✔
231	dataloader = DataLoader(dataset, batch_size=self._fd_batch_size, shuffle=False)	1✔
232	task_fisher_diags = fd_compute(	1✔
233	self._model,
234	self._train_forward,
235	dataloader, # type: ignore
236	self.device,
237	self._criterion,
238	)
239	# Update the fisher diagonals buffer with the computed values
240	for i in range(len(self._fisher_diags)):	1✔
241	self._fisher_diags[i].mul_(self._gamma).add_(task_fisher_diags[i])	1✔
242
243	def _update_anchor_params(self) -> None:	1✔
244	"""Update anchored parameters to the current model weights."""
245	for param, anchor_param in zip(	1✔
246	self._model.parameters(), self._anchor_params, strict=True
247	):
248	anchor_param.copy_(param.detach())	1✔
249
250	def _test_forward(self, x: Tensor) -> Tensor:	1✔
251	"""Compute logits for inference, optionally applying a test-task mask."""
252	y_hat = self._model(x)	1✔
253	if self._task_mask is not None and self._mask_test:	1✔
254	y_hat = self._task_mask[self._test_task] * y_hat	1✔
255	return y_hat	1✔
256
257	def _train_forward(self, x: Tensor) -> Tensor:	1✔
258	"""Compute logits for training, optionally applying a train-task mask."""
259	y_hat = self._model(x)	1✔
260	if self._task_mask is not None and self._mask_train:	1✔
UNCOV 261	y_hat = self._task_mask[self._train_task] * y_hat	×
262	return y_hat	1✔
263
264	def _regularisation_loss(self) -> torch.Tensor:	1✔
265	"""Return the EWC regularisation loss for the current task."""
266	if self._train_task < 1:	1✔
267	return torch.tensor(0.0, device=self.device)	1✔
268	return weighted_l2_reg(	1✔
269	self._model.parameters(),
270	self._anchor_params,
271	self._fisher_diags,
272	device=self.device,
273	)

adaptive-machine-learning / CapyMOA / 25189796363

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous