(prototype) One SimCLR example, one benchmark, and the gate between them

benchmarks/bench/__init__.py

@@ -0,0 +1,10 @@
+"""What operates a benchmark run without touching the update.
+
+Loaders, logging and checkpoint plumbing live here. Anything that changes what
+the optimiser sees lives in the method's own benchmark file, where a reviewer
+reads it beside the equation.
+"""
+
+from benchmarks.bench.datamodule import ImageDataModule
+
+__all__ = ["ImageDataModule"]

benchmarks/bench/datamodule.py

@@ -0,0 +1,93 @@
+"""The data side of a benchmark: loaders, and nothing that touches the update."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any, Callable, List, Optional
+
+from pytorch_lightning import LightningDataModule
+from torch.utils.data import DataLoader
+
+from lightly.data import LightlyDataset
+from lightly.data.sample import collate
+from lightly.transforms.torchvision_v2_compatibility import torchvision_transforms as T
+
+
+class ImageDataModule(LightningDataModule):
+    """Train, kNN-train and validation loaders over two image folders.
+
+    The kNN loader reads the training images under the validation transform, so
+    the probe scores the encoder rather than the augmentations. It is the first
+    validation loader, and the validation set is the second; the probes read that
+    ordering.
+
+    Args:
+        train_dir: Folder of training images.
+        val_dir: Folder of validation images.
+        transform: The method's transform. It returns views.
+        batch_size: Per-device batch size.
+        size: Crop size the validation transform resizes to.
+        normalize: Channel means and standard deviations.
+        num_workers: Worker processes per loader.
+    """
+
+    def __init__(
+        self,
+        train_dir: Path,
+        val_dir: Path,
+        transform: Callable[..., Any],
+        batch_size: int,
+        size: int,
+        normalize: Any,
+        num_workers: int = 8,
+    ) -> None:
+        super().__init__()
+        self.train_dir = train_dir
+        self.val_dir = val_dir
+        self.transform = transform
+        self.batch_size = batch_size
+        self.num_workers = num_workers
+        self.val_transform = T.Compose(
+            [
+                T.Resize(int(size * 256 / 224)),
+                T.CenterCrop(size),
+                T.ToTensor(),
+                T.Normalize(mean=normalize["mean"], std=normalize["std"]),
+            ]
+        )
+
+    def _loader(
+        self,
+        directory: Path,
+        transform: Callable[..., Any],
+        shuffle: bool,
+        drop_last: bool,
+        collate_fn: Optional[Callable[..., Any]] = None,
+    ) -> DataLoader:
+        return DataLoader(
+            LightlyDataset(input_dir=str(directory), transform=transform),
+            batch_size=self.batch_size,
+            shuffle=shuffle,
+            drop_last=drop_last,
+            num_workers=self.num_workers,
+            persistent_workers=self.num_workers > 0,
+            collate_fn=collate_fn,
+        )
+
+    def train_dataloader(self) -> DataLoader:
+        return self._loader(
+            self.train_dir,
+            self.transform,
+            shuffle=True,
+            drop_last=True,
+            collate_fn=collate,
+        )
+
+    def val_dataloader(self) -> List[DataLoader]:
+        knn_train = self._loader(
+            self.train_dir, self.val_transform, shuffle=False, drop_last=False
+        )
+        validation = self._loader(
+            self.val_dir, self.val_transform, shuffle=False, drop_last=False
+        )
+        return [knn_train, validation]

benchmarks/bench/probes.py

@@ -0,0 +1,139 @@
+"""Probes a benchmark scores its encoder with.
+
+These take features and targets and return numbers. They do not inherit from
+``LightningModule``, do not log, and do not key their lifecycle on which
+dataloader a batch came from: the benchmark says when the bank is built and when
+it is scored, on two adjacent lines a reviewer can read.
+
+The public, method-agnostic version of this is ``lightly.eval``, which is a
+separate piece of work. Until it lands this is the smallest thing that serves the
+benchmarks in this tree.
+"""
+
+from __future__ import annotations
+
+from typing import Callable, Dict, List, Optional, Tuple
+
+import torch
+from torch import Tensor
+from torch.nn.functional import normalize
+
+from lightly.utils.benchmarking.knn import knn_predict
+from lightly.utils.benchmarking.topk import mean_topk_accuracy
+
+__all__ = ["KNNProbe"]
+
+
+class KNNProbe:
+    """Weighted kNN over a feature bank built from the training set.
+
+    Settings follow InstDisc, which is what the published benchmark numbers used.
+
+    Args:
+        num_classes: Classes in the dataset.
+        k: Neighbours that vote.
+        t: Temperature the similarities are reweighted with.
+        topk: Which top-k accuracies to report.
+        feature_dtype: Bank dtype. ``float16`` halves the memory.
+        normalize_features: Whether to L2-normalize before the search.
+    """
+
+    def __init__(
+        self,
+        num_classes: int,
+        k: int,
+        t: float,
+        topk: Tuple[int, ...] = (1, 5),
+        feature_dtype: torch.dtype = torch.float32,
+        normalize_features: bool = True,
+    ) -> None:
+        self.num_classes = num_classes
+        self.k = k
+        self.t = t
+        self.topk = topk
+        self.feature_dtype = feature_dtype
+        self.normalize_features = normalize_features
+        self._features: List[Tensor] = []
+        self._targets: List[Tensor] = []
+        self._bank: Optional[Tensor] = None
+        self._labels: Optional[Tensor] = None
+
+    def _prepare(self, features: Tensor) -> Tensor:
+        if self.normalize_features:
+            features = normalize(features, dim=1)
+        return features.to(self.feature_dtype)
+
+    def reset(self) -> None:
+        """Drops the bank, so the next epoch starts from nothing."""
+        self._features = []
+        self._targets = []
+        self._bank = None
+        self._labels = None
+
+    def add(self, features: Tensor, targets: Tensor) -> None:
+        """Adds one batch of training features to the bank.
+
+        Args:
+            features: Encoder features of shape ``(B, D)``.
+            targets: Labels of shape ``(B,)``.
+        """
+        self._features.append(self._prepare(features).cpu())
+        self._targets.append(targets.cpu())
+
+    def build(
+        self,
+        gather: Optional[Callable[[Tensor], Tensor]] = None,
+        device: Optional[torch.device] = None,
+    ) -> None:
+        """Closes the bank so it can be scored against.
+
+        Args:
+            gather:
+                Collects a tensor from every rank, returning
+                ``(world_size, B, ...)``. ``LightningModule.all_gather`` is one.
+                Omit it on a single rank.
+            device: Where the bank is scored. Defaults to where it was built.
+
+        Raises:
+            ValueError: If no features were added.
+        """
+        if not self._features:
+            raise ValueError("the kNN bank is empty: call add() before build()")
+        features = torch.cat(self._features, dim=0)
+        targets = torch.cat(self._targets, dim=0)
+        if gather is not None:
+            features = gather(features)
+            targets = gather(targets)
+        # (dim, world_size * batch) is the layout knn_predict reads.
+        self._bank = features.flatten(end_dim=-2).t().contiguous().to(device)
+        self._labels = targets.flatten().contiguous().to(device)
+        self._features = []
+        self._targets = []
+
+    def score(self, features: Tensor, targets: Tensor) -> Dict[str, Tensor]:
+        """Scores one batch against the bank.
+
+        Args:
+            features: Encoder features of shape ``(B, D)``.
+            targets: Labels of shape ``(B,)``.
+
+        Returns:
+            One entry per top-k, keyed ``val_knn_top{k}``.
+
+        Raises:
+            ValueError: If the bank has not been built.
+        """
+        if self._bank is None or self._labels is None:
+            raise ValueError("the kNN bank is not built: call build() before score()")
+        predicted = knn_predict(
+            feature=self._prepare(features),
+            feature_bank=self._bank,
+            feature_labels=self._labels,
+            num_classes=self.num_classes,
+            knn_k=self.k,
+            knn_t=self.t,
+        )
+        topk = mean_topk_accuracy(
+            predicted_classes=predicted, targets=targets, k=self.topk
+        )
+        return {f"val_knn_top{k}": accuracy for k, accuracy in topk.items()}

benchmarks/simclr/README.md

@@ -0,0 +1,43 @@
+# SimCLR
+
+```bash
+torchrun --nproc_per_node=8 -m benchmarks.simclr.benchmark \
+    --train-dir /datasets/imagenet/train --val-dir /datasets/imagenet/val
+```
+
+```bash
+python -m benchmarks.simclr.benchmark --dataset cifar10 \
+    --train-dir ./cifar/train --val-dir ./cifar/val --devices 1 --strategy auto
+```
+
+`datasets.py` holds one row per dataset. Each row states every field, so two rows
+read side by side show every difference there is. `benchmark.py` takes one row.
+
+## Numbers
+
+| Row | Backbone | Batch | Epochs | Linear top1 | kNN top1 | Produced by |
+| :-- | :-- | --: | --: | --: | --: | :-- |
+| imagenet | ResNet-50 | 4096 | 100 | — | — | not run yet |
+| cifar10 | ResNet-18 | 256 | 100 | — | — | not run yet |
+
+The number in the repository README — 63.2 linear, 73.9 kNN — came from
+`benchmarks/imagenet/resnet50/simclr.py`, last touched in c7928120, on the run
+logged as `imagenet_resnet50_simclr_2023-06-22_09-11-13`. That file is deleted
+here and the row above has not reproduced it. Three things differ, so treat the
+old number as a target rather than as this benchmark's result:
+
+- **Batch size and learning rate.** The old run used 256 with square-root
+  scaling. The `imagenet` row is the paper's 4096 at lr 4.8.
+- **The online probe.** It ran on both views' features out of the SSL forward and
+  its head rode the method's optimiser. Here it sees one view under `no_grad`, so
+  `val_online_cls_top*` moves. The SSL loss is unaffected: the probe's input was
+  already detached.
+- **The kNN probe.** Same settings, k=200 and t=0.1, but `bench/probes.py`
+  replaces the dataloader-index state machine with `reset`, `add`, `build` and
+  `score`, called from `validation_step`.
+
+## What is not here yet
+
+`bench/probes.py` is the smallest thing that serves this benchmark.
+`lightly.eval` is where a public, method-agnostic probe goes, and the online
+linear probe still lives inside `benchmark.py` until then.