Coverage for cafl4ds/ssl/base.py: 100%
44 statements
« prev ^ index » next coverage.py v7.15.0, created at 2026-07-17 14:39 +0000
1"""The SSL-method interface (the ``C`` flag) and shared init/checkpoint utilities.
3An :class:`SSLMethod` bundles the shared encoder with a self-supervised objective. The
4streaming loop only ever calls two things on it:
6* :meth:`SSLMethod.training_step` — compute a self-supervised loss on a batch of **raw
7 images** (no labels ever enter here); the loop backprops and steps the optimizer.
8* :meth:`SSLMethod.encode` — map images to the pooled backbone embedding the health
9 instruments and probes read (no gradient).
11The ``C`` factor of the experiment matrix selects the concrete method (:mod:`cafl4ds.ssl.mae`
12or :mod:`cafl4ds.ssl.simsiam`); the ``I`` factor selects the encoder initialization, applied
13here via :func:`apply_encoder_init` (``from_scratch`` leaves the random init in place;
14``pretrained`` loads a checkpoint produced by an *IID* pre-pass, so stream correlation never
15contaminates the starting point).
16"""
18from __future__ import annotations
20from abc import ABC, abstractmethod
21from pathlib import Path
23import torch
24from loguru import logger
25from torch import nn
27from cafl4ds.models.vit import TinyViTEncoder
30class SSLMethod(nn.Module, ABC): # type: ignore[misc] # nn.Module is Any without torch stubs (mypy hook env)
31 """A self-supervised method: a shared encoder plus a training objective."""
33 def __init__(self, encoder: TinyViTEncoder) -> None:
34 """Store the shared backbone encoder.
36 Args:
37 encoder: The :class:`~cafl4ds.models.vit.TinyViTEncoder` backbone under study.
38 """
39 super().__init__()
40 self.encoder = encoder
42 @property
43 @abstractmethod
44 def name(self) -> str:
45 """Short method identifier (e.g. ``"mae"``, ``"simsiam"``) for logging."""
47 @abstractmethod
48 def training_step(self, imgs: torch.Tensor) -> torch.Tensor:
49 """Compute the self-supervised loss on a batch of raw images.
51 Args:
52 imgs: A batch of images ``[B, C, H, W]``. Labels never enter this path.
54 Returns:
55 A scalar loss tensor with gradient, ready for ``.backward()``.
56 """
58 def per_sample_loss(self, imgs: torch.Tensor) -> torch.Tensor:
59 """Return the *per-sample* self-supervised loss ``[B]`` (no gradient).
61 Where :meth:`training_step` reduces to a single scalar, this keeps the loss for each
62 image separately — the free per-frame informativeness signal the plan calls out (MAE's
63 per-patch reconstruction error; SimSiam's per-view negative cosine). The loss-gate knob
64 (:class:`~cafl4ds.filters.loss_gate.LossGate`) reads it to keep the high-loss frames;
65 later phases reuse it for novelty scoring. It is a *selection heuristic*, computed under
66 ``no_grad`` and never backpropagated, so it is free to draw its own augmentation/mask.
68 Args:
69 imgs: A batch of raw images ``[B, C, H, W]``. Labels never enter this path.
71 Returns:
72 A detached ``[B]`` tensor of per-sample losses (same orientation as the scalar
73 :meth:`training_step`: lower is better).
75 Raises:
76 NotImplementedError: If the concrete method does not define a per-sample loss.
77 """
78 raise NotImplementedError(f"{type(self).__name__} does not define a per-sample loss.")
80 def encode(self, imgs: torch.Tensor) -> torch.Tensor:
81 """Return the pooled backbone embedding used by the instruments/probes.
83 Args:
84 imgs: A batch of images ``[B, C, H, W]``.
86 Returns:
87 The pooled embedding ``[B, embed_dim]`` (no gradient tracking).
88 """
89 with torch.no_grad():
90 return self.encoder.embed(imgs)
93def load_encoder_checkpoint(encoder: TinyViTEncoder, checkpoint: str | Path) -> None:
94 """Load encoder weights from a checkpoint saved by :func:`save_encoder_checkpoint`.
96 Args:
97 encoder: The encoder to load weights into (in place).
98 checkpoint: Path to a ``state_dict`` file for the encoder.
100 Raises:
101 FileNotFoundError: If ``checkpoint`` does not exist.
102 """
103 path = Path(checkpoint)
104 if not path.is_file():
105 raise FileNotFoundError(f"pretrained checkpoint not found: {path}")
106 state = torch.load(path, map_location="cpu")
107 encoder.load_state_dict(state)
108 logger.info(f"loaded pretrained encoder weights from {path}")
111def save_encoder_checkpoint(encoder: TinyViTEncoder, checkpoint: str | Path) -> None:
112 """Save an encoder's ``state_dict`` (produced by the IID pre-pass).
114 Args:
115 encoder: The encoder whose weights to save.
116 checkpoint: Destination path (parent directories are created).
117 """
118 path = Path(checkpoint)
119 path.parent.mkdir(parents=True, exist_ok=True)
120 # Move tensors to CPU before serializing: the checkpoint is device-agnostic on disk
121 # (loading uses ``map_location="cpu"``), and saving an ``hpu`` state_dict directly trips a
122 # Habana storage-copy bug. ``.contiguous()`` normalizes strides so the CPU copy is clean.
123 state = {k: v.detach().contiguous().cpu() for k, v in encoder.state_dict().items()}
124 torch.save(state, path)
125 logger.info(f"saved pretrained encoder weights to {path}")
128def apply_encoder_init(
129 encoder: TinyViTEncoder, mode: str = "from_scratch", checkpoint: str | Path | None = None
130) -> None:
131 """Apply the ``I`` (initialization) factor to the encoder.
133 Args:
134 encoder: The encoder to initialize.
135 mode: ``"from_scratch"`` (keep the random init) or ``"pretrained"`` (warm-start from
136 ``checkpoint``).
137 checkpoint: Path to the warm-start checkpoint; required when ``mode`` is
138 ``"pretrained"``.
140 Raises:
141 ValueError: If ``mode`` is unknown, or ``"pretrained"`` without a ``checkpoint``.
142 """
143 if mode == "from_scratch":
144 logger.info("encoder init: from_scratch (random weights)")
145 return
146 if mode == "pretrained":
147 if checkpoint is None:
148 raise ValueError("init mode 'pretrained' requires a checkpoint path.")
149 load_encoder_checkpoint(encoder, checkpoint)
150 return
151 raise ValueError(f"unknown init mode {mode!r}; expected 'from_scratch' or 'pretrained'.")