Coverage for cafl4ds/filters/loss_gate.py: 100%

1"""Loss-gate — keep the high-loss frames (SOFed 2022; loss-based importance). 

2 

3A loss-based importance knob: score each incoming frame by the SSL loss the **current** model 

4incurs on it (:meth:`~cafl4ds.ssl.base.SSLMethod.per_sample_loss` — MAE's reconstruction error, 

5SimSiam's negative cosine) and keep only the fraction with the highest loss. High loss = the 

6model reconstructs / matches this frame poorly, i.e. it is the informative, not-yet-learned 

7part of the stream; low-loss frames are already-fit and cheap to skip. This is the "closest 

8prior art" loss-importance selection (baseline **B3**), run here as a streaming, label-free, 

9single-pass training-selection signal. 

10 

11Like SemDeDup this is an *admission* knob (it narrows the batch) and consults only the live 

12model, never labels or a stored gradient. 

13""" 

14 

15from __future__ import annotations 

16 

17import torch 

18 

19from cafl4ds.data.streams import StreamBatch 

20from cafl4ds.filters.base import Filter, FilterContext 

21 

22 

23class LossGate(Filter): 

24 """Admission knob: keep the top ``keep_frac`` of frames by current per-sample SSL loss.""" 

25 

26 def __init__(self, keep_frac: float = 0.5) -> None: 

27 """Configure the loss-gate. 

28 

29 Args: 

30 keep_frac: Fraction of each batch to keep (the highest-loss frames), in ``(0, 1]``. 

31 At least one frame is always kept. ``1.0`` keeps the whole batch (a no-op gate). 

32 

33 Raises: 

34 ValueError: If ``keep_frac`` is not in ``(0, 1]``. 

35 """ 

36 if not 0.0 < keep_frac <= 1.0: 

37 raise ValueError(f"keep_frac must be in (0, 1]; got {keep_frac}.") 

38 self.keep_frac = keep_frac 

39 

40 def select(self, batch: StreamBatch, ctx: FilterContext) -> torch.Tensor: 

41 """Return the highest-loss ``keep_frac`` subset of the batch (in stream order). 

42 

43 Args: 

44 batch: The incoming label-free stream batch. 

45 ctx: The live filter context; its ``method`` scores per-frame loss (fast edge). 

46 

47 Returns: 

48 The kept images ``[K, C, H, W]`` with ``K = max(1, round(keep_frac * B))``. 

49 """ 

50 b = batch.images.shape[0] 

51 keep_n = max(1, round(self.keep_frac * b)) 

52 if keep_n >= b: 

53 return batch.images 

54 losses = ctx.method.per_sample_loss(batch.images) # [B], no grad, higher = keep 

55 top = torch.topk(losses, keep_n).indices 

56 top, _ = torch.sort(top) # preserve stream order among the kept frames 

57 return batch.images[top]