lcc.classifiers

 49class BaseClassifier(pl.LightningModule):
 50    """
 51    Base image classifier class that supports LCC.
 52
 53    Warning:
 54        When subclassing this, remember that the forward method must be able to
 55        deal with either `Tensor` or `Batch` inputs, and must return a logit
 56        `Tensor`.
 57    """
 58
 59    lcc_data: dict[str, LatentClusteringData] | None = None
 60    """
 61    If LCC is applied, then this is non `None` and updated at the begining of
 62    each epoch. See also `full_dataset_latent_clustering`.
 63    """
 64
 65    standard_loss: nn.Module
 66    """'Standard' loss to use together with LCC."""
 67
 68    accuracy_top1: nn.Module
 69    """Top-1 accuracy metric."""
 70
 71    accuracy_top5: nn.Module | None = None
 72    """Top-5 accuracy metric."""
 73
 74    def __init__(
 75        self,
 76        n_classes: int,
 77        lcc_submodules: list[str] | None = None,
 78        lcc_kwargs: dict | None = None,
 79        ce_weight: float = 1,
 80        image_key: Any = 0,
 81        label_key: Any = 1,
 82        **kwargs: Any,
 83    ) -> None:
 84        """
 85        Args:
 86            n_classes (int):
 87            lcc_submodules (list[str] | None, optional): Submodules to consider
 88                for the latent correction loss. If `None` or `[]`, LCC is not
 89                performed
 90            lcc_kwargs (dict | None, optional): Optional parameters for LCC.
 91                Expected entries (all optional) are:
 92                * **weight (float):** Defaults to $10^{-4}$
 93                * **class_selection
 94                    (`lcc.correction.LCCClassSelection` | None):** Defaults to
 95                    `None`, which means all classes are considered for
 96                    correction
 97                * **interval (int):** Apply LCC every `interval` epochs.
 98                    Defaults to $1$, meaning LCC will be applied every epoch
 99                    (after warmup).
100                * **warmup (int):** Number of epochs to wait before
101                    starting LCC. Defaults to $0$, meaning LCC will start
102                    immediately.
103                * **k (int):** Number of nearest neighbors to consider for LCC,
104                  and Louvain clustering.
105                * **pca_dim (int):** Samples are reduced to this dimension
106                  before constructing the KNN graph. This must be at most the
107                  batch size.
108                * **loss (`"exact"` or `"randomized"`)**: Way the LCC loss is
109                  computed from the clustering data. See `lcc.correction.loss`.
110                * **ccspc (int)**: If the loss type is `"randomized"`, then this
111                  parameter specified the number of CC samples per cluster to
112                  keep as potential correction targets. See also
113                  `lcc.correction.loss.RandomizedLCCLoss`.
114                * **clustering_method (`"louvain"` or `"peer_pressure"`)**:
115                  Clustering algorithm for the dataset of latent
116                  representations.
117            ce_weight (float, optional): Weight of the cross-entropy loss in the
118                clustering-CE loss. Ignored if LCC is not applied. Defaults to
119                $1$.
120            image_key (Any, optional): A batch passed to the model can be a
121                tuple (most common) or a dict. This parameter specifies the key
122                to use to retrieve the input tensor.
123            label_key (Any, optional): Analogous to `image_key`.
124            kwargs: Forwarded to
125                [`pl.LightningModule`](https://lightning.ai/docs/pytorch/stable/common/lightning_module.html#)
126        """
127        super().__init__(**kwargs)
128        self.save_hyperparameters(
129            "ce_weight",
130            "image_key",
131            "label_key",
132            "lcc_kwargs",
133            "lcc_submodules",
134            "n_classes",
135        )
136        if lcc_submodules:
137            validate_lcc_kwargs(lcc_kwargs)
138        self.standard_loss = torch.nn.CrossEntropyLoss()
139        acc_kw: dict[str, Any] = {
140            "task": "multiclass",
141            "num_classes": n_classes,
142            "average": "micro",
143        }
144        self.accuracy_top1 = tm.Accuracy(top_k=1, **acc_kw)  # type: ignore
145        if n_classes > 5:
146            self.accuracy_top5 = tm.Accuracy(top_k=5, **acc_kw)  # type: ignore
147
148    def _evaluate(self, batch: Batch, stage: str | None = None) -> Tensor:
149        """Self-explanatory"""
150        image_key = self.hparams["image_key"]
151        label_key = self.hparams["label_key"]
152        x, y = batch[image_key], batch[label_key].to(self.device)
153        latent: dict[str, Tensor] = {}
154        logits = self.forward_intermediate(
155            x, self.lcc_submodules, latent, keep_gradients=True
156        )
157        assert isinstance(logits, Tensor)
158        loss_ce = self.standard_loss(logits, y)
159        if self.lcc_data and stage == "train":
160            idx = to_array(batch["_idx"])
161            _losses = [
162                self.lcc_data[sm].loss(
163                    z, y_true=y, y_clst=self.lcc_data[sm].y_clst[idx]
164                )
165                for sm, z in latent.items()
166            ]
167            loss_lcc = (
168                torch.stack(_losses).mean()
169                if _losses
170                else torch.tensor(0.0, requires_grad=True)
171                # ↑ actually need grad?
172            )
173            lcc_weight = self.hparams.get("lcc_kwargs", {}).get("weight", 1e-4)
174            loss = self.hparams["ce_weight"] * loss_ce + lcc_weight * loss_lcc
175            self.log(f"{stage}/lcc", loss_lcc, sync_dist=True)
176        else:
177            loss = loss_ce
178        if stage:
179            d = {
180                f"{stage}/loss": loss,
181                f"{stage}/ce": loss_ce,
182            }
183            if self.accuracy_top5 is not None:
184                d[f"{stage}/acc5"] = self.accuracy_top5(logits, y)
185            self.log_dict(d, sync_dist=True)
186            self.log(
187                f"{stage}/acc",
188                self.accuracy_top1(logits, y),
189                prog_bar=True,
190                sync_dist=True,
191            )
192        return loss  # type: ignore
193
194    def configure_optimizers(self) -> Any:
195        optimizer = torch.optim.SGD(self.parameters(), lr=1e-1, momentum=0.9)
196        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
197            optimizer, T_max=50
198        )
199        return {
200            "optimizer": optimizer,
201            "lr_scheduler": {"scheduler": scheduler},
202        }
203
204    def forward_intermediate(
205        self,
206        inputs: Tensor | Batch | list[Tensor] | Sequence[Batch],
207        submodules: list[str],
208        output_dict: dict,
209        keep_gradients: bool = False,
210    ) -> Tensor | list[Tensor]:
211        """
212        Runs the model and collects the output of specified submodules. The
213        intermediate outputs are stored in `output_dict` under the
214        corresponding submodule name. In particular, this method has side
215        effects.
216
217        Args:
218            x (Tensor | Batch | list[Tensor] | list[Batch]): If batched (i.e.
219                `x` is a list), then so is the output of this function and the
220                entries in the `output_dict`
221            submodules (list[str]):
222            output_dict (dict):
223            keep_gradients (bool, optional): If `True`, the tensors in
224                `output_dict` keep their gradients (if they had some on the
225                first place). If `False`, they are detached and moved to the
226                CPU.
227        """
228
229        def maybe_detach(x: Tensor) -> Tensor:
230            return x if keep_gradients else x.detach().cpu()
231
232        def create_hook(key: str) -> Callable[[nn.Module, Any, Any], None]:
233            def hook(_model: nn.Module, _args: Any, out: Any) -> None:
234                if (
235                    isinstance(out, (list, tuple))
236                    and len(out) == 1
237                    and isinstance(out[0], Tensor)
238                ):
239                    out = out[0]
240                elif (  # Special case for ViTs
241                    isinstance(out, (list, tuple))
242                    and len(out) == 2
243                    and isinstance(out[0], Tensor)
244                    and not isinstance(out[1], Tensor)
245                ):
246                    out = out[0]
247                elif not isinstance(out, Tensor):
248                    raise ValueError(
249                        f"Unsupported latent object type: {type(out)}: {out}."
250                    )
251                if batched:
252                    if key not in output_dict:
253                        output_dict[key] = []
254                    output_dict[key].append(maybe_detach(out))
255                else:
256                    output_dict[key] = maybe_detach(out)
257
258            return hook
259
260        batched = isinstance(inputs, (list, tuple))
261        handles: list[RemovableHandle] = []
262        for name in submodules:
263            submodule = self.get_submodule(name)
264            handles.append(submodule.register_forward_hook(create_hook(name)))
265        if batched:
266            logits = [
267                maybe_detach(
268                    self.forward(
269                        batch
270                        if isinstance(batch, Tensor)
271                        else batch[self.hparams["image_key"]]
272                    )
273                )
274                for batch in inputs
275            ]
276        else:
277            logits = maybe_detach(  # type: ignore
278                self.forward(
279                    inputs
280                    if isinstance(inputs, Tensor)
281                    else inputs[self.hparams["image_key"]]
282                )
283            )
284        for h in handles:
285            h.remove()
286        return logits
287
288    @staticmethod
289    def get_image_processor(model_name: str, **kwargs: Any) -> Callable:
290        """
291        Returns an image processor for the model. By defaults, returns the
292        identity function.
293        """
294        return lambda input: input
295
296    @property
297    def lcc_submodules(self) -> list[str]:
298        """
299        Returns the list of submodules considered for LCC, whith correct prefix
300        if needed.
301        """
302        return self.hparams.get("lcc_submodules") or []
303
304    def on_train_batch_end(self, *args: Any, **kwargs: Any) -> None:
305        """Just logs all optimizer's learning rate"""
306        log_optimizers_lr(self, sync_dist=True)
307        super().on_train_batch_end(*args, **kwargs)
308
309    def on_train_epoch_end(self) -> None:
310        """Cleans up training specific temporary attributes"""
311        self.lcc_data = None
312        super().on_train_epoch_end()
313
314    def on_train_epoch_start(self) -> None:
315        """
316        Performs dataset-wide latent clustering and stores the results in
317        private attribute `BaseClassifier._lc_data`.
318        """
319        # wether to apply LCC this epoch
320        lcc_kwargs = self.hparams.get("lcc_kwargs") or {}
321        do_lcc = (
322            # we are passed warmup (lcc_warmup being None is equivalent to no
323            # warmup)...
324            self.current_epoch >= (lcc_kwargs.get("warmup") or 0)
325            and (
326                # ... and an LCC interval is specified...
327                lcc_kwargs.get("interval") is not None
328                # ... and the current epoch can have LCC done...
329                and self.current_epoch % int(lcc_kwargs.get("interval", 1))
330                == 0
331            )
332            # ... and there are submodule selected for LCC...
333            and self.lcc_submodules
334            # ... and the LCC weight is non-zero
335            and lcc_kwargs.get("weight", 0) > 0
336        )
337        if do_lcc:
338            # The import has to be here to prevent circular imports while having
339            # all the FDLC logic in a separate file
340            from .fdlc import full_dataset_latent_clustering
341
342            with temporary_directory(self) as tmp_path:
343                self.lcc_data = full_dataset_latent_clustering(
344                    model=self,
345                    output_dir=tmp_path,
346                    tqdm_style="console",
347                )
348        super().on_train_epoch_start()
349
350    def on_train_start(self) -> None:
351        """
352        Explicitly registers hyperparameters and metrics. You'd think Lightning
353        would do this automatically, but nope.
354        """
355        self.logger.log_hyperparams(  # type: ignore
356            self.hparams,
357            {
358                s + "/" + m: np.nan
359                for s, m in product(
360                    ["train", "val"], ["acc", "loss", "ce", "lcc"]
361                )
362            },
363        )
364        super().on_train_start()
365
366    def test_step(self, batch: Batch, *_: Any, **__: Any) -> Tensor:
367        """Override from `pl.LightningModule.test_step`."""
368        return self._evaluate(batch, "test")
369
370    def training_step(self, batch: Batch, *_: Any, **__: Any) -> Tensor:
371        """Override from `pl.LightningModule.training_step`."""
372        return self._evaluate(batch, "train")
373
374    def validation_step(self, batch: Batch, *_: Any, **__: Any) -> Tensor:
375        """Override from `pl.LightningModule.validation_step`."""
376        return self._evaluate(batch, "val")

26def get_classifier_cls(model_name: str) -> type[BaseClassifier]:
27    """
28    Returns the classifier class to use for a given model name.
29
30    Args:
31        model_name (str):
32    """
33    if model_name.startswith("timm/"):
34        return TimmClassifier
35    if "/" in model_name:
36        return HuggingFaceClassifier
37    return TorchvisionClassifier

11class HuggingFaceClassifier(WrappedClassifier):
12    """
13    Pretrained classifier model loaded from the [HuggingFace model
14    hub](https://huggingface.co/models?pipeline_tag=image-classification).
15    """
16
17    def __init__(
18        self,
19        model_name: str,
20        n_classes: int,
21        head_name: str | None = None,
22        **kwargs: Any,
23    ) -> None:
24        """
25        See also:
26            `lcc.classifiers.WrappedClassifier.__init__` and
27            `lcc.classifiers.BaseClassifier.__init__`.
28
29        Args:
30            model_name (str): Model name as in the [HuggingFace model
31                hub](https://huggingface.co/models?pipeline_tag=image-classification).
32                If the model name starts with `timm/`, use
33                `lcc.classifiers.TimmClassifier` instead.
34            n_classes (int): See `lcc.classifiers.WrappedClassifier.__init__`.
35            head_name (str | None, optional): See
36                `lcc.classifiers.WrappedClassifier.__init__`.
37        """
38        if model_name.startswith("timm/"):
39            raise ValueError(
40                "If the model name starts with `timm/`, use "
41                "`lcc.classifiers.TimmClassifier` instead."
42            )
43        model = AutoModelForImageClassification.from_pretrained(model_name)
44        super().__init__(model, n_classes, head_name, **kwargs)
45        self.save_hyperparameters()
46
47    @staticmethod
48    def get_image_processor(
49        model_name: str, **__: Any
50    ) -> Callable[[dict[str, Any]], dict[str, Any]]:
51        """
52        Wraps the HuggingFace `AutoImageProcessor` associated to a given model.
53
54        Args:
55            model_name (str): Model name as in the [HuggingFace model
56                hub](https://huggingface.co/models?pipeline_tag=image-classification).
57                Must not start with `timm/`.
58
59        Returns:
60            A callable that uses a
61            [`transformers.AutoImageProcessor`](https://huggingface.co/docs/transformers/v4.44.2/en/model_doc/auto#transformers.AutoImageProcessor)
62            under the hood.
63        """
64        hf_transorm = AutoImageProcessor.from_pretrained(model_name)
65
66        def _transform(batch: dict[str, Any]) -> dict[str, Any]:
67            return {
68                k: (
69                    hf_transorm(
70                        [img.convert("RGB") for img in v], return_tensors="pt"
71                    )["pixel_values"]
72                    # TODO: pass image_key from DS ↓
73                    if k in ["img", "image", "jpg", "png"]
74                    else v
75                )
76                for k, v in batch.items()
77            }
78
79        return _transform

30@dataclass
31class LatentClusteringData:
32    """
33    Convenience struct that holds some latent clustering correction data for a
34    given latent space.
35    """
36
37    loss: LCCLoss
38    """
39    The actual LCC loss object (which is callable) that compute the LCC loss for
40    LCC correction.
41    """
42
43    # TODO: Write the y_clst to disk and loading alongside the train dataset, so
44    # that this whole datastructure becomes unnecessary
45    y_clst: np.ndarray
46    """`(N,)` vector of cluster labels."""

12class TimmClassifier(WrappedClassifier):
13    """
14    Pretrained classifier model loaded from the [HuggingFace model
15    hub](https://huggingface.co/models?pipeline_tag=image-classification)
16    uploaded by the `timm` team.
17    """
18
19    image_transform: Callable
20
21    def __init__(
22        self,
23        model_name: str,
24        n_classes: int,
25        head_name: str | None = None,
26        pretrained: bool = True,
27        **kwargs: Any,
28    ) -> None:
29        """
30        See also:
31            `lcc.classifiers.WrappedClassifier.__init__` and
32            `lcc.classifiers.BaseClassifier.__init__`.
33
34        Args:
35            model_name (str): Model name as in the [HuggingFace model
36                hub](https://huggingface.co/models?pipeline_tag=image-classification).
37                Must start with `timm/`.
38            n_classes (int): See `lcc.classifiers.WrappedClassifier.__init__`.
39            head_name (str | None, optional): See
40                `lcc.classifiers.WrappedClassifier.__init__`.
41            pretrained (bool, optional): Defaults to `True`.
42        """
43        if not model_name.startswith("timm/"):
44            raise ValueError(
45                "The model isn't a timm model (its name does not start with "
46                "`timm/`). Use `lcc.classifiers.HuggingFaceClassifier` "
47                "instead."
48            )
49        model = timm.create_model(model_name, pretrained=pretrained)
50        super().__init__(model, n_classes, head_name, **kwargs)
51        self.save_hyperparameters()
52
53    @staticmethod
54    def get_image_processor(model_name: str, **kwargs: Any) -> Callable:
55        """
56        Wraps the HuggingFace `AutoImageProcessor` associated to a given model.
57
58        See also:
59            [`timm.create_transform`](https://huggingface.co/docs/timm/reference/data#timm.data.create_transform).
60
61        Args:
62            model_name (str): Model name as in the [HuggingFace model
63                hub](https://huggingface.co/models?pipeline_tag=image-classification).
64                Must start with `timm/`.
65
66        Returns:
67            A callable that uses a [Torchvision
68            transform](https://pytorch.org/vision/0.19/transforms.html) under
69            the hood.
70        """
71        model = timm.create_model(model_name, pretrained=False)
72        conf = timm.data.resolve_model_data_config(model)
73        conf["is_training"], conf["no_aug"] = True, True
74        timm_transform = timm.data.create_transform(**conf)
75
76        def _transform(batch: dict[str, Any]) -> dict[str, Any]:
77            return {
78                k: (
79                    (
80                        [timm_transform(img.convert("RGB")) for img in v]
81                        if isinstance(v, list)
82                        else timm_transform(v)
83                    )
84                    # TODO: pass image_key from DS ↓
85                    if k in ["img", "image", "jpg", "png"]
86                    else v
87                )
88                for k, v in batch.items()
89            }
90
91        return _transform
92
93    def forward(self, inputs: Tensor | Batch, *_: Any, **__: Any) -> Tensor:
94        x: Tensor = (
95            inputs if isinstance(inputs, Tensor) else inputs[self.image_key]
96        )
97        return self.model(x.to(self.device))  # type: ignore

12class TorchvisionClassifier(WrappedClassifier):
13    """
14    A torchvision classifier wrapped as a `WrappedClassifier`. See
15    https://pytorch.org/vision/stable/models.html#classification .
16    """
17
18    def __init__(
19        self,
20        model_name: str,
21        n_classes: int,
22        head_name: str | None = None,
23        weights: Any = "DEFAULT",
24        **kwargs: Any,
25    ) -> None:
26        model = get_model(model_name, weights=weights)
27        super().__init__(model, n_classes, head_name, **kwargs)
28        self.save_hyperparameters()
29
30    @staticmethod
31    def get_image_processor(
32        model_name: str, weights: str = "DEFAULT", **__: Any
33    ) -> Callable[[dict[str, Any]], dict[str, Any]]:
34        """
35        Creates an image processor based on the transform object of the model's
36        chosen weights. For example,
37
38            TorchvisionClassifier.get_image_processor("alexnet")
39
40        is analogous to
41
42            get_model_weights("alexnet")["DEFAULT"].transforms()
43        """
44
45        transform = tr.Compose(
46            [tr.RGB(), get_model_weights(model_name)[weights].transforms()]
47        )
48
49        def _processor(batch: dict[str, Any]) -> dict[str, Any]:
50            return {
51                k: (
52                    [transform(img) for img in v]
53                    # TODO: pass image_key from DS ↓
54                    if k in ["img", "image", "jpg", "png"]
55                    else v
56                )
57                for k, v in batch.items()
58            }
59
60        return _processor