Coverage for cafl4ds/loop.py: 93%

1"""The streaming SSL adaptation loop (Phase 0). 

2 

3Ties the pieces together: pull batches from the :class:`~cafl4ds.data.streams.EraStream`, pass 

4each through the :class:`~cafl4ds.filters.base.Filter` (Phase 0: the B-floor accept-all knob), 

5run one self-supervised update on the accepted images, log the loss every step, and every 

6``eval_every`` steps read the :class:`~cafl4ds.monitor.HealthMonitor` and log the health 

7series. Loss and health land in the *same* run log as *separate* series so they can be read 

8side by side over steps (the Phase-0 exit criterion). 

9 

10By default this is a single-pass streaming loop: no replay, no shuffling of the incoming 

11order — the model sees the correlated stream exactly as it arrives. It also supports a 

12**multi-epoch** regime (``epochs > 1``, the stream re-iterated in the same order) with an 

13optional per-step **LR scheduler** — the fair training regime the P0.2.1 healthy-baseline 

14calibration needs (the correlated single pass is still the default, and the Phase-1 setting). 

15""" 

16 

17from __future__ import annotations 

18 

19from collections.abc import Iterable, Sized 

20 

21import torch 

22from loguru import logger 

23from torch import optim 

24from torch.optim.lr_scheduler import LRScheduler 

25 

26from cafl4ds.data.streams import StreamBatch 

27from cafl4ds.eval import PerEraProbe 

28from cafl4ds.filters.base import Filter, FilterContext 

29from cafl4ds.monitor import HealthMonitor 

30from cafl4ds.run_log import RunLogger 

31from cafl4ds.ssl.base import SSLMethod 

32 

33# BatchNorm heads (SimSiam) and per-patch stats need at least two samples in a batch. 

34_MIN_BATCH = 2 

35 

36 

37class StreamingLoop: 

38 """Runs the streaming SSL adaptation loop with a single selection filter.""" 

39 

40 def __init__( 

41 self, 

42 stream: Iterable[StreamBatch], 

43 method: SSLMethod, 

44 optimizer: optim.Optimizer, 

45 selection_filter: Filter, 

46 monitor: HealthMonitor, 

47 run_logger: RunLogger, 

48 eval_every: int = 5, 

49 epochs: int = 1, 

50 scheduler: LRScheduler | None = None, 

51 grad_clip: float | None = 1.0, 

52 device: str = "cpu", 

53 era_evaluator: PerEraProbe | None = None, 

54 ) -> None: 

55 """Build the loop. 

56 

57 Args: 

58 stream: An iterable of :class:`~cafl4ds.data.streams.StreamBatch` (the ``F`` factor). 

59 method: The SSL method to adapt (the ``C`` factor). 

60 optimizer: Optimizer over ``method.parameters()``. 

61 selection_filter: The selection knob (the ``A`` factor); Phase 0 uses B-floor. 

62 monitor: The health monitor read every ``eval_every`` (global) steps. 

63 run_logger: The run log receiving the loss and health series. 

64 eval_every: Run the monitor every this many steps (and always at the end). 

65 epochs: Number of passes over the stream. ``1`` (default) is the single-pass 

66 streaming setting; ``>1`` re-iterates the same stream in the same order (the 

67 multi-epoch calibration regime). Steps are numbered **globally** across epochs. 

68 scheduler: Optional LR scheduler stepped once per optimizer step (e.g. 

69 :func:`~cafl4ds.schedule.warmup_cosine_schedule`), or ``None`` for a flat LR. 

70 grad_clip: Global grad-norm clip value, or ``None`` to disable. 

71 device: Torch device to run on (``"cpu"`` for Phase-0 smoke; ``"hpu"`` later). 

72 era_evaluator: Optional probe-on-past evaluator (the *validating* axis). When set, 

73 the current encoder is probed over all seen eras at each era boundary and at the 

74 end, building the accuracy matrix behind Backward Transfer / Forgetting. ``None`` 

75 (default) runs no downstream probing — the loop is unchanged. 

76 """ 

77 self.stream = stream 

78 self.method = method 

79 self.optimizer = optimizer 

80 self.selection_filter = selection_filter 

81 self.monitor = monitor 

82 self.run_logger = run_logger 

83 self.eval_every = eval_every 

84 self.epochs = epochs 

85 self.scheduler = scheduler 

86 self.grad_clip = grad_clip 

87 self.device = torch.device(device) 

88 self.era_evaluator = era_evaluator 

89 

90 def run(self) -> RunLogger: 

91 """Run the stream (``epochs`` passes), logging loss and health over global steps. 

92 

93 Returns: 

94 The run logger (closed), for convenience. 

95 """ 

96 self.method.to(self.device) 

97 # Global step numbering across epochs. batch.step restarts at 0 each pass, so offset 

98 # by a fixed batches-per-epoch stride to keep steps monotonic (and the eval cadence / 

99 # drift reference consistent across a multi-epoch run). A single pass needs no offset. 

100 if self.epochs > 1 and not isinstance(self.stream, Sized): 100 ↛ 101line 100 didn't jump to line 101 because the condition on line 100 was never true

101 raise TypeError("multi-epoch loop (epochs > 1) requires a stream with a known length.") 

102 batches_per_epoch = len(self.stream) if isinstance(self.stream, Sized) else 0 

103 last_step, last_era = -1, 0 

104 prev_era: int | None = None 

105 for epoch in range(self.epochs): 

106 for batch in self.stream: 

107 step = epoch * batches_per_epoch + batch.step 

108 if prev_era is not None and batch.era != prev_era: 

109 self._probe_past(prev_era) # the era just ended — record its probe-on-past row 

110 prev_era = batch.era 

111 last_era = batch.era 

112 moved = StreamBatch(images=batch.images.to(self.device), era=batch.era, step=step) 

113 accepted = self.selection_filter.select(moved, FilterContext(method=self.method, step=step)) 

114 if accepted.shape[0] < _MIN_BATCH: 114 ↛ 115line 114 didn't jump to line 115 because the condition on line 114 was never true

115 logger.debug(f"step {step}: skipping batch of {accepted.shape[0]} (< {_MIN_BATCH}).") 

116 continue 

117 loss = self._update(accepted) 

118 self.run_logger.log_loss(step, batch.era, loss) 

119 if step % self.eval_every == 0: 

120 self.run_logger.log_health(step, batch.era, self.monitor.measure(self.method, step)) 

121 last_step = step 

122 if last_step >= 0 and last_step % self.eval_every != 0: # always end on a health reading 

123 self.run_logger.log_health(last_step, last_era, self.monitor.measure(self.method, last_step)) 

124 if last_step >= 0: 124 ↛ 126line 124 didn't jump to line 126 because the condition on line 124 was always true

125 self._probe_past(last_era) # final row: current encoder over every era seen 

126 logger.info("streaming loop complete\n" + self.run_logger.tabulate()) 

127 self.run_logger.close() 

128 return self.run_logger 

129 

130 def _probe_past(self, era: int) -> None: 

131 """Record one probe-on-past row for the just-finished ``era`` (if an evaluator is set). 

132 

133 Probes in eval mode (BatchNorm/dropout off, as the health monitor does), then restores 

134 the training mode so the loop is unperturbed. 

135 

136 Args: 

137 era: Index of the era that just completed; the encoder is scored on eras ``0..era``. 

138 """ 

139 if self.era_evaluator is None: 

140 return 

141 was_training = self.method.training 

142 self.method.eval() 

143 try: 

144 self.era_evaluator.record(self.method.encode, era) 

145 finally: 

146 self.method.train(was_training) 

147 

148 def _update(self, images: torch.Tensor) -> float: 

149 """Run one SSL optimization step on the accepted images. 

150 

151 Args: 

152 images: The accepted image batch ``[K, C, H, W]``. 

153 

154 Returns: 

155 The scalar loss value for this step. 

156 """ 

157 self.method.train() 

158 self.optimizer.zero_grad() 

159 loss = self.method.training_step(images) 

160 loss.backward() 

161 if self.grad_clip is not None: 161 ↛ 163line 161 didn't jump to line 163 because the condition on line 161 was always true

162 torch.nn.utils.clip_grad_norm_(self.method.parameters(), self.grad_clip) 

163 self.optimizer.step() 

164 if self.scheduler is not None: 

165 self.scheduler.step() 

166 return float(loss.item())