25192586337

Committed 30 Apr 2026 10:31PM UTC coverage: 75.462% (+0.6%) from 74.909%

Build # 25192586337

Build Type

Pull #355

github

Committed by

web-flow

Commit Message

Merge 287330aab into 2ef7fe73e

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

Coverage Stats

341 of 347 new or added lines in 18 files covered. (98.27%)

6 existing lines in 4 files now uncovered.

7427 of 9842 relevant lines covered (75.46%)

0.75 hits per line

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

99.05

/src/capymoa/ocl/util/_replay.py

from typing import Dict, Sequence, OrderedDict, Tuple
from typing_extensions import override
from torch import Tensor, nn
from torch.utils.data import TensorDataset
from abc import abstractmethod, ABC
from capymoa.ocl.util._buffer import BufferDict
import torch


def _detach(**batch: Dict[str, Tensor]) -> Dict[str, Tensor]:
    """Detach a batch of tensors from the computation graph."""
    return {key: tensor.detach() for key, tensor in batch.items()}


class ReplayBuffer(ABC, nn.Module):
    @abstractmethod
    def update(self, **batch: Dict[str, Tensor]) -> None:
        """Update the replay buffer with new examples.

        :param batch: Dictionary of tensors representing a batch of examples.
        """
        ...

    def sample(self, n: int) -> OrderedDict[str, Tensor]:
        """Sample ``n`` examples from the replay buffer.

        :param n: Number of examples to sample
        :return: Dictionary of tensors representing the sampled examples.
        """
        indices = torch.randint(0, self.count, (n,))
        return {key: buffer[indices] for key, buffer in self._buffer.items()}

    def array(self) -> Dict[str, Tensor]:
        """Return the replay buffer as a dictionary of tensors."""
        return {key: buffer[: self.count] for key, buffer in self._buffer.items()}

    def dataset_view(self) -> TensorDataset:
        """Return a TensorDataset view of the replay buffer."""
        return TensorDataset(*self.array().values())

    @property
    def capacity(self) -> int:
        """Return the maximum number of samples that can be stored in the coreset."""
        return self._capacity

    @property
    def count(self) -> int:
        """Return the current number of samples in the coreset."""
        assert self._count <= self._capacity
        return self._count

    @property
    def device(self) -> torch.device:
        return next(iter(self._buffer.values())).device

    def __init__(
        self,
        capacity: int,
        buffers: Dict[str, Tuple[Sequence[int], torch.dtype]],
        rng: torch.Generator = torch.default_generator,
    ) -> None:
        super().__init__()
        self._capacity = capacity
        self._rng = rng
        self._count = 0
        self._i = 0
        self._buffer = BufferDict(
            {
                key: torch.zeros((capacity, *shape), dtype=dtype)
                for key, (shape, dtype) in buffers.items()
            }
        )

    @classmethod
    def new_xybuffer(
        cls,
        capacity: int,
        x_shape: Sequence[int],
        rng: torch.Generator = torch.default_generator,
    ) -> "ReplayBuffer":
        """Standard buffer with features ``x: (n, *x_shape) f32`` and labels ``y: (n,) i64``."""
        return cls(
            capacity, {"x": (x_shape, torch.float32), "y": ((), torch.long)}, rng
        )


class ReservoirSampler(ReplayBuffer):
    @override
    def update(self, **batch: Dict[str, Tensor]) -> None:
        assert set(batch.keys()) == set(self._buffer.keys())
        batch = _detach(**batch)
        batch_size = next(iter(batch.values())).shape[0]
        for key, values in batch.items():
            assert values.shape[0] == batch_size
            batch[key] = values.to(self._buffer[key].device)

        for i in range(batch_size):
            if self.count < self.capacity:
                # Fill the reservoir.
                for key, values in batch.items():
                    self._buffer[key][self.count] = values[i]
                self._count += 1
            else:
                # Standard reservoir sampling replacement.
                index = int(
                    torch.randint(0, self._i + 1, (1,), generator=self._rng).item()
                )
                if index < self.capacity:
                    for key, values in batch.items():
                        self._buffer[key][index] = values[i]
            self._i += 1


class GreedySampler(ReplayBuffer):
    """Update the buffer with every new example, replacing a random example from the
    majority class if the buffer is full.
    """

    @override
    def update(self, **batch: Dict[str, Tensor]) -> None:
        assert set(batch.keys()) == set(self._buffer.keys())
        assert "y" in batch
        batch = _detach(**batch)

        batch_size = next(iter(batch.values())).shape[0]
        prepared_batch: Dict[str, Tensor] = {}
        for key, values in batch.items():
            assert values.shape[0] == batch_size
            prepared_batch[key] = values.to(self._buffer[key].device)

        for i in range(batch_size):
            if self.count < self.capacity:
                # Room left in the coreset for this example.
                for key, values in prepared_batch.items():
                    self._buffer[key][self.count] = values[i]
                self._count += 1
                continue

            # Coreset is full, replace a random example from the majority class.
            y_buffer = self._buffer["y"][: self.count]
            classes, counts = y_buffer.unique(return_counts=True)
            replace_class = classes[counts.argmax()]
            candidate_indices = (y_buffer == replace_class).nonzero(as_tuple=True)[0]
            idx = int(
                torch.randint(
                    0, len(candidate_indices), (1,), generator=self._rng
                ).item()
            )
            replace_idx = int(candidate_indices[idx].item())
            for key, values in prepared_batch.items():
                self._buffer[key][replace_idx] = values[i]


class SlidingWindow(ReplayBuffer):
    """Update the buffer with every new example, replacing the oldest example if the
    buffer is full.
    """

    @override
    def update(self, **batch: Dict[str, Tensor]) -> None:
        assert set(batch.keys()) == set(self._buffer.keys())
        batch = _detach(**batch)

        batch_size = next(iter(batch.values())).shape[0]
        prepared_batch: Dict[str, Tensor] = {}
        for key, values in batch.items():
            assert values.shape[0] == batch_size
            prepared_batch[key] = values.to(self._buffer[key].device)

        # Calculate where the batch ends
        end_idx = self._i + batch_size

        if end_idx <= self.capacity:
            # Case 1: Simple slice (no wrap-around)
            for key, values in prepared_batch.items():
                self._buffer[key][self._i : end_idx] = values
        else:
            # Case 2: Wrap-around (split the batch)
            mid_point = self.capacity - self._i

            # Fill the end of the buffer
            for key, values in prepared_batch.items():
                self._buffer[key][self._i :] = values[:mid_point]

            # Wrap the remainder to the start
            wrap_size = batch_size - mid_point
            for key, values in prepared_batch.items():
                self._buffer[key][:wrap_size] = values[mid_point:]

        # Update index and count
        self._i = (self._i + batch_size) % self.capacity
        self._count = min(self._count + batch_size, self.capacity)

1	from typing import Dict, Sequence, OrderedDict, Tuple	1✔
2	from typing_extensions import override	1✔
3	from torch import Tensor, nn	1✔
4	from torch.utils.data import TensorDataset	1✔
5	from abc import abstractmethod, ABC	1✔
6	from capymoa.ocl.util._buffer import BufferDict	1✔
7	import torch	1✔
8
9
10	def _detach(**batch: Dict[str, Tensor]) -> Dict[str, Tensor]:	1✔
11	"""Detach a batch of tensors from the computation graph."""
12	return {key: tensor.detach() for key, tensor in batch.items()}	1✔
13
14
15	class ReplayBuffer(ABC, nn.Module):	1✔
16	@abstractmethod	1✔
17	def update(self, **batch: Dict[str, Tensor]) -> None:	1✔
18	"""Update the replay buffer with new examples.
19
20	:param batch: Dictionary of tensors representing a batch of examples.
21	"""
22	...
23
24	def sample(self, n: int) -> OrderedDict[str, Tensor]:	1✔
25	"""Sample ``n`` examples from the replay buffer.
26
27	:param n: Number of examples to sample
28	:return: Dictionary of tensors representing the sampled examples.
29	"""
30	indices = torch.randint(0, self.count, (n,))	1✔
31	return {key: buffer[indices] for key, buffer in self._buffer.items()}	1✔
32
33	def array(self) -> Dict[str, Tensor]:	1✔
34	"""Return the replay buffer as a dictionary of tensors."""
35	return {key: buffer[: self.count] for key, buffer in self._buffer.items()}	1✔
36
37	def dataset_view(self) -> TensorDataset:	1✔
38	"""Return a TensorDataset view of the replay buffer."""
39	return TensorDataset(*self.array().values())	1✔
40
41	@property	1✔
42	def capacity(self) -> int:	1✔
43	"""Return the maximum number of samples that can be stored in the coreset."""
44	return self._capacity	1✔
45
46	@property	1✔
47	def count(self) -> int:	1✔
48	"""Return the current number of samples in the coreset."""
49	assert self._count <= self._capacity	1✔
50	return self._count	1✔
51
52	@property	1✔
53	def device(self) -> torch.device:	1✔
NEW UNCOV 54	return next(iter(self._buffer.values())).device	×
55
56	def __init__(	1✔
57	self,
58	capacity: int,
59	buffers: Dict[str, Tuple[Sequence[int], torch.dtype]],
60	rng: torch.Generator = torch.default_generator,
61	) -> None:
62	super().__init__()	1✔
63	self._capacity = capacity	1✔
64	self._rng = rng	1✔
65	self._count = 0	1✔
66	self._i = 0	1✔
67	self._buffer = BufferDict(	1✔
68	{
69	key: torch.zeros((capacity, *shape), dtype=dtype)
70	for key, (shape, dtype) in buffers.items()
71	}
72	)
73
74	@classmethod	1✔
75	def new_xybuffer(	1✔
76	cls,
77	capacity: int,
78	x_shape: Sequence[int],
79	rng: torch.Generator = torch.default_generator,
80	) -> "ReplayBuffer":
81	"""Standard buffer with features ``x: (n, *x_shape) f32`` and labels ``y: (n,) i64``."""
82	return cls(	1✔
83	capacity, {"x": (x_shape, torch.float32), "y": ((), torch.long)}, rng
84	)
85
86
87	class ReservoirSampler(ReplayBuffer):	1✔
88	@override	1✔
89	def update(self, **batch: Dict[str, Tensor]) -> None:	1✔
90	assert set(batch.keys()) == set(self._buffer.keys())	1✔
91	batch = _detach(**batch)	1✔
92	batch_size = next(iter(batch.values())).shape[0]	1✔
93	for key, values in batch.items():	1✔
94	assert values.shape[0] == batch_size	1✔
95	batch[key] = values.to(self._buffer[key].device)	1✔
96
97	for i in range(batch_size):	1✔
98	if self.count < self.capacity:	1✔
99	# Fill the reservoir.
100	for key, values in batch.items():	1✔
101	self._buffer[key][self.count] = values[i]	1✔
102	self._count += 1	1✔
103	else:
104	# Standard reservoir sampling replacement.
105	index = int(	1✔
106	torch.randint(0, self._i + 1, (1,), generator=self._rng).item()
107	)
108	if index < self.capacity:	1✔
109	for key, values in batch.items():	1✔
110	self._buffer[key][index] = values[i]	1✔
111	self._i += 1	1✔
112
113
114	class GreedySampler(ReplayBuffer):	1✔
115	"""Update the buffer with every new example, replacing a random example from the
116	majority class if the buffer is full.
117	"""
118
119	@override	1✔
120	def update(self, **batch: Dict[str, Tensor]) -> None:	1✔
121	assert set(batch.keys()) == set(self._buffer.keys())	1✔
122	assert "y" in batch	1✔
123	batch = _detach(**batch)	1✔
124
125	batch_size = next(iter(batch.values())).shape[0]	1✔
126	prepared_batch: Dict[str, Tensor] = {}	1✔
127	for key, values in batch.items():	1✔
128	assert values.shape[0] == batch_size	1✔
129	prepared_batch[key] = values.to(self._buffer[key].device)	1✔
130
131	for i in range(batch_size):	1✔
132	if self.count < self.capacity:	1✔
133	# Room left in the coreset for this example.
134	for key, values in prepared_batch.items():	1✔
135	self._buffer[key][self.count] = values[i]	1✔
136	self._count += 1	1✔
137	continue	1✔
138
139	# Coreset is full, replace a random example from the majority class.
140	y_buffer = self._buffer["y"][: self.count]	1✔
141	classes, counts = y_buffer.unique(return_counts=True)	1✔
142	replace_class = classes[counts.argmax()]	1✔
143	candidate_indices = (y_buffer == replace_class).nonzero(as_tuple=True)[0]	1✔
144	idx = int(	1✔
145	torch.randint(
146	0, len(candidate_indices), (1,), generator=self._rng
147	).item()
148	)
149	replace_idx = int(candidate_indices[idx].item())	1✔
150	for key, values in prepared_batch.items():	1✔
151	self._buffer[key][replace_idx] = values[i]	1✔
152
153
154	class SlidingWindow(ReplayBuffer):	1✔
155	"""Update the buffer with every new example, replacing the oldest example if the
156	buffer is full.
157	"""
158
159	@override	1✔
160	def update(self, **batch: Dict[str, Tensor]) -> None:	1✔
161	assert set(batch.keys()) == set(self._buffer.keys())	1✔
162	batch = _detach(**batch)	1✔
163
164	batch_size = next(iter(batch.values())).shape[0]	1✔
165	prepared_batch: Dict[str, Tensor] = {}	1✔
166	for key, values in batch.items():	1✔
167	assert values.shape[0] == batch_size	1✔
168	prepared_batch[key] = values.to(self._buffer[key].device)	1✔
169
170	# Calculate where the batch ends
171	end_idx = self._i + batch_size	1✔
172
173	if end_idx <= self.capacity:	1✔
174	# Case 1: Simple slice (no wrap-around)
175	for key, values in prepared_batch.items():	1✔
176	self._buffer[key][self._i : end_idx] = values	1✔
177	else:
178	# Case 2: Wrap-around (split the batch)
179	mid_point = self.capacity - self._i	1✔
180
181	# Fill the end of the buffer
182	for key, values in prepared_batch.items():	1✔
183	self._buffer[key][self._i :] = values[:mid_point]	1✔
184
185	# Wrap the remainder to the start
186	wrap_size = batch_size - mid_point	1✔
187	for key, values in prepared_batch.items():	1✔
188	self._buffer[key][:wrap_size] = values[mid_point:]	1✔
189
190	# Update index and count
191	self._i = (self._i + batch_size) % self.capacity	1✔
192	self._count = min(self._count + batch_size, self.capacity)	1✔

adaptive-machine-learning / CapyMOA / 25192586337

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