Coverage for cafl4ds/data/streams.py: 99%
111 statements
« prev ^ index » next coverage.py v7.15.0, created at 2026-07-17 14:39 +0000
1"""Streaming datasets: order a :class:`~cafl4ds.data.sources.DataSource` into eras.
3The stream is the ``F`` (data-order) factor of the experiment. Phase 0 uses a single
4:class:`EraStream` with two orderings:
6* ``order="class_blocked"`` — the **synthetic correlation** of Phase 1a: all of class 0, then
7 all of class 1, …. Each class block is an **era**. This is the correlated stream whose
8 degradation the instruments must catch.
9* ``order="iid"`` — a shuffled single pass (one era). This is the *well-behaved reference*
10 used to produce the ``I=pretrained`` warm-start checkpoint — correlation must **not** leak
11 into the clean starting point.
13Labels are used only to (a) build the class-blocked ordering and (b) construct held-out eval
14sets; they never appear in a :class:`StreamBatch`, so they can never enter the SSL update.
15Each stream reserves, per class, a disjoint held-out **probe support**, **probe query**, and
16**per-era** eval set before ordering the remaining images into the training stream.
17"""
19from __future__ import annotations
21from collections.abc import Iterator
22from dataclasses import dataclass, field
24import torch
26from cafl4ds.data.sources import DataSource
29@dataclass(frozen=True)
30class StreamBatch:
31 """One batch delivered by a stream — images only, tagged with its era and step."""
33 images: torch.Tensor
34 """Image batch ``[B, C, H, W]``. Contains no labels, by construction."""
35 era: int
36 """Index of the era (class block) this batch belongs to."""
37 step: int
38 """Global step index of this batch within the single-pass stream."""
41@dataclass(frozen=True)
42class EvalSet:
43 """A held-out labeled evaluation set. Labels are used HERE ONLY (probes/eval)."""
45 images: torch.Tensor
46 """Images ``[M, C, H, W]``."""
47 labels: torch.Tensor
48 """Integer labels ``[M]`` — used only by the probes, never by training."""
51@dataclass(frozen=True)
52class EvalSets:
53 """The held-out eval sets a stream exposes to the health monitor."""
55 probe_support: EvalSet
56 """Balanced held-out set the probes fit on (kNN database / linear-head training)."""
57 probe_query: EvalSet
58 """Balanced held-out set the probes score on; also the *fixed reference* whose embeddings
59 the drift instruments track across checkpoints."""
60 per_era: dict[int, EvalSet] = field(default_factory=dict)
61 """Per-era (per-class) held-out sets, reserved for later per-era forgetting measures."""
64class EraStream:
65 """A single-pass stream that orders a data source into eras.
67 Reserves disjoint per-class held-out eval sets, then orders the remaining images either
68 as class blocks (``"class_blocked"``) or shuffled (``"iid"``), and delivers them as
69 label-free :class:`StreamBatch` batches.
70 """
72 def __init__(
73 self,
74 source: DataSource,
75 batch_size: int = 32,
76 order: str = "class_blocked",
77 class_order: list[int] | None = None,
78 support_per_class: int = 20,
79 query_per_class: int = 20,
80 era_eval_per_class: int = 10,
81 max_train_per_class: int | None = None,
82 drop_last: bool = False,
83 seed: int = 0,
84 ) -> None:
85 """Build the stream (loads the source and constructs the splits eagerly).
87 Args:
88 source: The data source to order.
89 batch_size: Number of images per delivered batch.
90 order: ``"class_blocked"`` (eras = class blocks) or ``"iid"`` (one shuffled era).
91 class_order: Explicit class ordering for ``"class_blocked"``; defaults to
92 ``0, 1, …, num_classes - 1``.
93 support_per_class: Images per class reserved for the probe support set.
94 query_per_class: Images per class reserved for the probe query / drift set.
95 era_eval_per_class: Images per class reserved for the per-era eval set.
96 max_train_per_class: If set, cap the *training* images per class after the
97 held-out reservations (keeps smoke runs short).
98 drop_last: Whether to drop a final short batch.
99 seed: RNG seed for the held-out sampling and IID shuffle.
101 Raises:
102 ValueError: If ``order`` is unknown, or a class has too few images to satisfy the
103 requested held-out reservations.
104 """
105 if order not in ("class_blocked", "iid"):
106 raise ValueError(f"unknown order {order!r}; expected 'class_blocked' or 'iid'.")
107 self.batch_size = batch_size
108 self.order = order
109 self.drop_last = drop_last
110 self._generator = torch.Generator().manual_seed(seed)
112 images, labels = source.load()
113 self._images = images
114 classes = sorted(set(labels.tolist()))
115 self._class_order = class_order if class_order is not None else classes
117 train_indices_by_era: list[tuple[int, torch.Tensor]] = []
118 support_idx, query_idx = [], []
119 per_era_eval: dict[int, EvalSet] = {}
120 for era, cls in enumerate(self._class_order):
121 cls_idx = (labels == cls).nonzero(as_tuple=True)[0]
122 perm = cls_idx[torch.randperm(cls_idx.numel(), generator=self._generator)]
123 need = support_per_class + query_per_class + era_eval_per_class
124 if perm.numel() <= need:
125 raise ValueError(
126 f"class {cls} has {perm.numel()} images but {need} are reserved for eval; "
127 "reduce the per-class reservations or use more data."
128 )
129 s, q, e = support_per_class, query_per_class, era_eval_per_class
130 support_idx.append(perm[:s])
131 query_idx.append(perm[s : s + q])
132 era_eval = perm[s + q : s + q + e]
133 train = perm[s + q + e :]
134 if max_train_per_class is not None:
135 train = train[:max_train_per_class]
136 per_era_eval[era] = EvalSet(images[era_eval], labels[era_eval])
137 train_indices_by_era.append((era, train))
139 self._eval_sets = EvalSets(
140 probe_support=self._make_eval(torch.cat(support_idx), labels),
141 probe_query=self._make_eval(torch.cat(query_idx), labels),
142 per_era=per_era_eval,
143 )
144 self._order_stream = self._build_order(train_indices_by_era)
146 def _make_eval(self, idx: torch.Tensor, labels: torch.Tensor) -> EvalSet:
147 """Materialize an :class:`EvalSet` from an index tensor.
149 Args:
150 idx: Indices into the loaded images/labels.
151 labels: The full label vector.
153 Returns:
154 The corresponding held-out :class:`EvalSet`.
155 """
156 return EvalSet(self._images[idx], labels[idx])
158 def _build_order(self, train_by_era: list[tuple[int, torch.Tensor]]) -> list[tuple[int, int]]:
159 """Flatten per-era training indices into an ordered ``(era, image_index)`` stream.
161 Args:
162 train_by_era: ``(era, indices)`` pairs of training images per era.
164 Returns:
165 A list of ``(era, image_index)`` in the stream's delivery order.
166 """
167 ordered: list[tuple[int, int]] = []
168 for era, idx in train_by_era:
169 for i in idx.tolist():
170 ordered.append((era, i))
171 if self.order == "iid":
172 # IID is a single era: shuffle globally and relabel every item to era 0 so batches
173 # are not fragmented at (now meaningless) class boundaries.
174 perm = torch.randperm(len(ordered), generator=self._generator).tolist()
175 ordered = [(0, ordered[i][1]) for i in perm]
176 return ordered
178 @property
179 def eval_sets(self) -> EvalSets:
180 """The held-out eval sets (probe support/query and per-era)."""
181 return self._eval_sets
183 @property
184 def num_eras(self) -> int:
185 """Number of eras in the stream (class blocks, or 1 for IID)."""
186 return 1 if self.order == "iid" else len(self._class_order)
188 def __len__(self) -> int:
189 """Number of batches the stream will deliver.
191 Mirrors :meth:`__iter__` exactly, including the rule that **a batch is never split
192 across eras**: at every era boundary the current (possibly partial) batch is flushed.
193 For ``class_blocked`` that means each class block contributes ``ceil(block / batch)``
194 batches — strictly more than ``ceil(total / batch)`` whenever a block does not divide
195 evenly. Only the very last partial batch of the whole stream honours ``drop_last``.
196 (For ``iid`` there is a single era, so this reduces to ``ceil(total / batch_size)``.)
197 A correct count is load-bearing for the multi-epoch loop: it is the per-epoch stride
198 used to number global steps and to size the LR schedule's horizon.
199 """
200 count = 0
201 buffer = 0
202 buffer_era: int | None = None
203 for era, _ in self._order_stream:
204 if buffer_era is not None and era != buffer_era and buffer:
205 count += 1 # era-boundary flush (partial batch), always emitted
206 buffer = 0
207 buffer_era = era
208 buffer += 1
209 if buffer == self.batch_size:
210 count += 1
211 buffer = 0
212 if buffer and not self.drop_last and buffer_era is not None: 212 ↛ 214line 212 didn't jump to line 214 because the condition on line 212 was always true
213 count += 1 # final partial batch — the only one drop_last affects
214 return count
216 def __iter__(self) -> Iterator[StreamBatch]:
217 """Yield label-free :class:`StreamBatch` batches in stream order.
219 A batch is never split across eras: each batch carries a single era id (the era of
220 its first sample), and a new batch is started at every era boundary.
222 Yields:
223 Successive :class:`StreamBatch` batches over the single-pass stream.
224 """
225 step = 0
226 buffer_idx: list[int] = []
227 buffer_era: int | None = None
228 for era, image_index in self._order_stream:
229 if buffer_era is not None and era != buffer_era and buffer_idx:
230 yield self._make_batch(buffer_idx, buffer_era, step)
231 step += 1
232 buffer_idx = []
233 buffer_era = era
234 buffer_idx.append(image_index)
235 if len(buffer_idx) == self.batch_size:
236 yield self._make_batch(buffer_idx, era, step)
237 step += 1
238 buffer_idx = []
239 if buffer_idx and not self.drop_last and buffer_era is not None: 239 ↛ exitline 239 didn't return from function '__iter__' because the condition on line 239 was always true
240 yield self._make_batch(buffer_idx, buffer_era, step)
242 def _make_batch(self, idx: list[int], era: int, step: int) -> StreamBatch:
243 """Assemble a :class:`StreamBatch` from buffered indices.
245 Args:
246 idx: Image indices for this batch.
247 era: Era id for this batch.
248 step: Global step index.
250 Returns:
251 The assembled label-free batch.
252 """
253 return StreamBatch(images=self._images[torch.tensor(idx, dtype=torch.long)], era=era, step=step)