diff --git a/benchmarks/__init__.py b/benchmarks/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/benchmarks/bench/__init__.py b/benchmarks/bench/__init__.py new file mode 100644 index 000000000..f27fa99c1 --- /dev/null +++ b/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"] diff --git a/benchmarks/bench/datamodule.py b/benchmarks/bench/datamodule.py new file mode 100644 index 000000000..014db40ce --- /dev/null +++ b/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] diff --git a/benchmarks/bench/probes.py b/benchmarks/bench/probes.py new file mode 100644 index 000000000..dba3af744 --- /dev/null +++ b/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()} diff --git a/benchmarks/simclr/README.md b/benchmarks/simclr/README.md new file mode 100644 index 000000000..4a4a3a631 --- /dev/null +++ b/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. diff --git a/benchmarks/simclr/__init__.py b/benchmarks/simclr/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/benchmarks/simclr/benchmark.py b/benchmarks/simclr/benchmark.py new file mode 100644 index 000000000..6b9725963 --- /dev/null +++ b/benchmarks/simclr/benchmark.py @@ -0,0 +1,248 @@ +"""SimCLR on Lightning, parameterised by one row of datasets.py. + +The same method as ``examples/simclr.py``, written for the other job. The example +is the file to read; this one has to survive a week on a preemptible cluster, so +it hands the loop, the checkpoint and the step counter to Lightning. What it adds +on top of the example is the schedule, the probes, the logger, bf16, a checkpoint +and the dataset axis. + +The two files restate each other on purpose, and ``tests/test_simclr_agrees.py`` +is what holds them to the same method. They do not share numbers: the example is +one small configuration, the ``imagenet`` row is the paper's. + +Run it with:: + + torchrun --nproc_per_node=8 -m benchmarks.simclr.benchmark \\ + --train-dir /datasets/imagenet/train --val-dir /datasets/imagenet/val + + python -m benchmarks.simclr.benchmark --dataset cifar10 \\ + --train-dir ./cifar/train --val-dir ./cifar/val +""" + +from __future__ import annotations + +from argparse import ArgumentParser, Namespace +from pathlib import Path +from typing import Any, Dict, List, Tuple + +import torch +from pytorch_lightning import LightningModule, Trainer, seed_everything +from pytorch_lightning.callbacks import LearningRateMonitor, ModelCheckpoint +from pytorch_lightning.loggers import TensorBoardLogger +from torch import Tensor, no_grad +from torch.optim import SGD, Optimizer +from torchvision.models import resnet18, resnet50 + +from benchmarks.bench import ImageDataModule +from benchmarks.bench.probes import KNNProbe +from benchmarks.simclr.datasets import SETTINGS, Setting +from lightly.backbones import TorchvisionResNetBackbone, small_image_stem +from lightly.data.sample import Sample +from lightly.loss import NTXentLoss +from lightly.models.modules import SimCLRProjectionHead +from lightly.optim import LARS, CosineWarmupScheduler, param_groups +from lightly.transforms import SimCLRTransform +from lightly.utils.benchmarking import OnlineLinearClassifier + +BACKBONES = {"resnet18": resnet18, "resnet50": resnet50} +OPTIMIZERS = {"lars": LARS, "sgd": SGD} + + +def backbone(setting: Setting) -> TorchvisionResNetBackbone: + """Builds the row's backbone. Restated from the example; the gate compares it.""" + net = BACKBONES[setting.backbone]() + return TorchvisionResNetBackbone( + small_image_stem(net) if setting.stem == "small_image" else net + ) + + +def transform(setting: Setting) -> SimCLRTransform: + """Builds the row's transform. Restated from the example; the gate compares it.""" + return SimCLRTransform( + input_size=setting.size, + cj_strength=setting.color_jitter_strength, + gaussian_blur=setting.blur, + normalize=setting.normalize, + ) + + +class SimCLR(LightningModule): + """SimCLR, for one row. + + Args: + setting: The row from ``datasets.py`` to run. + """ + + def __init__(self, setting: Setting) -> None: + super().__init__() + self.save_hyperparameters({"setting": setting.name}) + self.setting = setting + self.backbone = backbone(setting) + self.head = SimCLRProjectionHead( + setting.feature_dim, setting.hidden_dim, setting.out_dim + ) + self.criterion = NTXentLoss( + temperature=setting.temperature, gather_distributed=True + ) + + # The probes stay inside the module until eval/ is public. Their optimiser + # is this one, which is what the published numbers were produced with. + self.online_classifier = OnlineLinearClassifier( + feature_dim=setting.feature_dim, num_classes=setting.num_classes + ) + self.knn_probe = KNNProbe( + num_classes=setting.num_classes, k=setting.knn_k, t=setting.knn_t + ) + + def forward(self, sample: Sample) -> Tensor: + # Both views go through the backbone and the head in one call, so every + # BatchNorm sees 2N samples. Running one view at a time gives it N, which + # is a different optimisation problem: the gradients of the two agree at + # cosine similarity 0.0837. + images = torch.cat([view.data for view in sample.views]) + z0, z1 = self.head(self.backbone.embed(images)).chunk(len(sample.views)) + loss: Tensor = self.criterion(z0, z1) + return loss + + def training_step(self, sample: Sample, batch_idx: int) -> Tensor: + loss = self(sample) + targets = sample.meta["target"] + self.log( + "train_loss", loss, prog_bar=True, sync_dist=True, batch_size=len(targets) + ) + + # The probe reads the encoder, not the projection, so it takes its own + # forward. One view under no_grad, where v1 reused both views from the + # SSL pass, so val_online_cls_top* moves against the published number. + with no_grad(): + features = self.backbone.embed(sample.views[0].data) + cls_loss, cls_log = self.online_classifier.training_step( + (features, targets), batch_idx + ) + self.log_dict(cls_log, sync_dist=True, batch_size=len(targets)) + return loss + cls_loss + + def validation_step( + self, + batch: Tuple[Tensor, Tensor, List[str]], + batch_idx: int, + dataloader_idx: int, + ) -> Tensor | None: + images, targets = batch[0], batch[1] + with no_grad(): + features = self.backbone.embed(images) + + if dataloader_idx == 0: + # Loader 0 is the training set under the validation transform: the + # kNN bank. Loader 1 is the validation set, scored against it. + if batch_idx == 0: + self.knn_probe.reset() + self.knn_probe.add(features, targets) + return None + + if batch_idx == 0: + self.knn_probe.build(gather=self.all_gather, device=self.device) + knn_log = self.knn_probe.score(features, targets) + cls_loss, cls_log = self.online_classifier.validation_step( + (features, targets), batch_idx + ) + self.log_dict( + {**knn_log, **cls_log}, + prog_bar=True, + sync_dist=True, + batch_size=len(targets), + ) + return cls_loss + + def on_validation_epoch_end(self) -> None: + self.knn_probe.reset() + + def configure_optimizers(self) -> Tuple[List[Optimizer], List[Dict[str, Any]]]: + setting = self.setting + groups = param_groups( + self.backbone, self.head, weight_decay=setting.weight_decay + ) + groups.append( + { + "name": "online_classifier", + "params": list(self.online_classifier.parameters()), + "weight_decay": 0.0, + } + ) + optimizer = OPTIMIZERS[setting.optimizer]( + groups, lr=setting.lr, momentum=setting.momentum + ) + steps_per_epoch = int( + self.trainer.estimated_stepping_batches / self.trainer.max_epochs + ) + scheduler = { + "scheduler": CosineWarmupScheduler( + optimizer=optimizer, + warmup_epochs=steps_per_epoch * setting.warmup_epochs, + max_epochs=int(self.trainer.estimated_stepping_batches), + ), + "interval": "step", + } + return [optimizer], [scheduler] + + +def parse_args() -> Namespace: + parser = ArgumentParser("SimCLR benchmark") + parser.add_argument( + "--dataset", type=str, default="imagenet", choices=list(SETTINGS) + ) + parser.add_argument("--train-dir", type=Path, default="/datasets/imagenet/train") + parser.add_argument("--val-dir", type=Path, default="/datasets/imagenet/val") + parser.add_argument("--out", type=Path, default=Path("benchmark_logs")) + parser.add_argument("--devices", type=int, default=-1) + parser.add_argument("--accelerator", type=str, default="gpu") + parser.add_argument( + "--strategy", type=str, default="ddp_find_unused_parameters_true" + ) + parser.add_argument("--num-workers", type=int, default=8) + parser.add_argument("--seed", type=int, default=11) + parser.add_argument("--fast-dev-run", action="store_true") + return parser.parse_args() + + +if __name__ == "__main__": + args = parse_args() + seed_everything(args.seed, workers=True) + setting = SETTINGS[args.dataset] + + # Setting.batch_size is the total across ranks, so the loaders get a share. + world_size = args.devices if args.devices > 0 else max(torch.cuda.device_count(), 1) + if setting.batch_size % world_size: + raise ValueError( + f"the {setting.name} row's batch size {setting.batch_size} does not " + f"divide into {world_size} devices" + ) + batch_size_per_device = setting.batch_size // world_size + + Trainer( + max_epochs=setting.epochs, + devices=args.devices, + accelerator=args.accelerator, + strategy=args.strategy, + precision="bf16-mixed", + default_root_dir=str(args.out), + logger=TensorBoardLogger(save_dir=str(args.out), name=setting.name), + callbacks=[ + ModelCheckpoint(save_last=True), + LearningRateMonitor(logging_interval="step"), + ], + fast_dev_run=args.fast_dev_run, + ).fit( + SimCLR(setting), + ImageDataModule( + train_dir=args.train_dir, + val_dir=args.val_dir, + transform=transform(setting), + batch_size=batch_size_per_device, + size=setting.size, + normalize=setting.normalize, + num_workers=args.num_workers, + ), + # Resume is one argument. + ckpt_path="last" if not args.fast_dev_run else None, + ) diff --git a/benchmarks/simclr/datasets.py b/benchmarks/simclr/datasets.py new file mode 100644 index 000000000..e14da0123 --- /dev/null +++ b/benchmarks/simclr/datasets.py @@ -0,0 +1,122 @@ +"""One row per dataset SimCLR is benchmarked on. + +Every row states every field. No defaults, no inheritance, no merging, and +nothing computed from anything else: two rows read side by side show every +difference there is, and a field left out is a type error rather than a silent +fallback. That is what keeps this a table rather than a configuration system, +and it is why ``lr`` is a literal instead of a scaling rule. + +Nothing here ships. ``benchmarks/`` is not in the wheel. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Dict, List + +from lightly.transforms.utils import CIFAR10_NORMALIZE, IMAGENET_NORMALIZE + + +@dataclass(frozen=True) +class Setting: + """One dataset SimCLR is benchmarked on. + + Attributes: + name: The row's key in ``SETTINGS``. + size: Crop size in pixels. + normalize: Channel means and standard deviations. + stem: ``standard`` or ``small_image``, the 3x3 stride-1 ResNet stem. + blur: Probability of Gaussian blur. + color_jitter_strength: Multiplies the four colour-jitter strengths. + backbone: ``resnet18`` or ``resnet50``. + feature_dim: The backbone's output width, and the head's input. + hidden_dim: The projection head's hidden width. + out_dim: The projection head's output width. + batch_size: Total batch size, summed over ranks. + epochs: Pretraining epochs. + temperature: NT-Xent temperature. + weight_decay: Applied to everything except norms and biases. + optimizer: ``lars`` or ``sgd``. + lr: The learning rate at this row's ``batch_size``. A literal, not a rule. + momentum: Optimiser momentum. + warmup_epochs: Linear warmup before the cosine decay. + num_classes: Classes in the dataset, for the probes. + knn_k: Neighbours the kNN probe votes over. + knn_t: Temperature the kNN probe reweights similarities with. + """ + + name: str + size: int + normalize: Dict[str, List[float]] + stem: str + blur: float + color_jitter_strength: float + backbone: str + feature_dim: int + hidden_dim: int + out_dim: int + batch_size: int + epochs: int + temperature: float + weight_decay: float + optimizer: str + lr: float + momentum: float + warmup_epochs: int + num_classes: int + knn_k: int + knn_t: float + + +# Chen et al. 2020, table 6 and appendix B.1. lr is the paper's linear rule at +# this batch size: 0.3 * 4096 / 256. +IMAGENET = Setting( + name="imagenet", + size=224, + normalize=IMAGENET_NORMALIZE, + stem="standard", + blur=0.5, + color_jitter_strength=1.0, + backbone="resnet50", + feature_dim=2048, + hidden_dim=2048, + out_dim=128, + batch_size=4096, + epochs=100, + temperature=0.1, + weight_decay=1e-6, + optimizer="lars", + lr=4.8, + momentum=0.9, + warmup_epochs=10, + num_classes=1000, + knn_k=200, + knn_t=0.1, +) + +# The settings examples/simclr.py is written out with. Small enough for one GPU. +CIFAR10 = Setting( + name="cifar10", + size=32, + normalize=CIFAR10_NORMALIZE, + stem="small_image", + blur=0.0, + color_jitter_strength=0.5, + backbone="resnet18", + feature_dim=512, + hidden_dim=512, + out_dim=128, + batch_size=256, + epochs=100, + temperature=0.5, + weight_decay=5e-4, + optimizer="sgd", + lr=0.06, + momentum=0.9, + warmup_epochs=0, + num_classes=10, + knn_k=200, + knn_t=0.1, +) + +SETTINGS = {setting.name: setting for setting in (IMAGENET, CIFAR10)} diff --git a/examples/simclr.py b/examples/simclr.py new file mode 100644 index 000000000..0f9a50c52 --- /dev/null +++ b/examples/simclr.py @@ -0,0 +1,96 @@ +"""SimCLR: a simple framework for contrastive learning of visual representations. + +Two augmented views of the same image are pulled together and pushed away from +every other image in the batch, with NT-Xent. Reference: Chen et al. 2020, +https://arxiv.org/abs/2002.05709. + +This file is one configuration, written out. It is ResNet-18 on CIFAR-10 at 32 +pixels, so it runs on one GPU and finishes; the paper's ImageNet settings are the +``imagenet`` row of ``benchmarks/simclr/datasets.py``, and that row is the one the +published number belongs to. The settings below are the ones this repository +benchmarked SimCLR on CIFAR-10 with. + +What changes for a different dataset is the block of constants and the four +lines that consume them. ``benchmarks/simclr/datasets.py`` is that change made +once per dataset, and ``tests/test_simclr_agrees.py`` is what holds the two files +to the same method. + +Run it against CIFAR-10 with:: + + python examples/simclr.py +""" + +import torch +from torch import Tensor +from torch.optim import SGD +from torch.utils.data import DataLoader +from torchvision.datasets import CIFAR10 +from torchvision.models import resnet18 + +from lightly.backbones import TorchvisionResNetBackbone, small_image_stem +from lightly.data.sample import Sample, collate +from lightly.loss import NTXentLoss +from lightly.models.modules import SimCLRProjectionHead +from lightly.optim import param_groups +from lightly.transforms import SimCLRTransform +from lightly.transforms.utils import CIFAR10_NORMALIZE + +# ResNet-18, CIFAR-10, 32 pixels. Small enough to run on one GPU. +BATCH_SIZE, EPOCHS, LR, MOMENTUM = 256, 100, 0.06, 0.9 +TEMPERATURE, WEIGHT_DECAY = 0.5, 5e-4 +INPUT_SIZE, FEATURE_DIM, HIDDEN_DIM, OUT_DIM = 32, 512, 512, 128 +COLOR_JITTER_STRENGTH, GAUSSIAN_BLUR, NORMALIZE = 0.5, 0.0, CIFAR10_NORMALIZE + +backbone = TorchvisionResNetBackbone(small_image_stem(resnet18())) +head = SimCLRProjectionHead(FEATURE_DIM, HIDDEN_DIM, OUT_DIM) +criterion = NTXentLoss(temperature=TEMPERATURE) +transform = SimCLRTransform( + input_size=INPUT_SIZE, + cj_strength=COLOR_JITTER_STRENGTH, + gaussian_blur=GAUSSIAN_BLUR, + normalize=NORMALIZE, +) + + +def forward(sample: Sample) -> Tensor: + # Both views go through the backbone and the head in one call, so every + # BatchNorm sees 2N samples. Running one view at a time gives it N, which is + # a different optimisation problem: the gradients of the two agree at cosine + # similarity 0.0837. + images = torch.cat([view.data for view in sample.views]) + z0, z1 = head(backbone.embed(images)).chunk(len(sample.views)) + return criterion(z0, z1) + + +if __name__ == "__main__": + device = "cuda" if torch.cuda.is_available() else "cpu" + backbone.to(device) + head.to(device) + + optimizer = SGD( + param_groups(backbone, head, weight_decay=WEIGHT_DECAY), + lr=LR, + momentum=MOMENTUM, + ) + dataset = CIFAR10("datasets/cifar10", download=True, transform=transform) + dataloader = DataLoader( + dataset, + batch_size=BATCH_SIZE, + shuffle=True, + drop_last=True, + num_workers=8, + collate_fn=collate, + ) + + print("Starting Training") + for epoch in range(EPOCHS): + total_loss = 0.0 + for sample in dataloader: + for view in sample.views: + view.data = view.data.to(device) + loss = forward(sample) + loss.backward() + optimizer.step() + optimizer.zero_grad() + total_loss += float(loss.detach()) + print(f"epoch: {epoch:>02}, loss: {total_loss / len(dataloader):.5f}") diff --git a/pyproject.toml b/pyproject.toml index ebbf26e7d..5573e92c6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -69,6 +69,8 @@ dev = [ "pandas", "toml", "torchmetrics", + # benchmarks/ logs to TensorBoard. Not needed by the wheel. + "tensorboard", # ruff and mypy should be the same version as defined in .pre-commit-config.yaml "ruff==0.12.7", # frozen version to avoid differences between CI and local dev machines "mypy==1.4.1", # frozen version to avoid differences between CI and local dev machines @@ -250,6 +252,18 @@ module = [ ] follow_imports = "skip" +# lightly.data.sample is typed, and the transforms depend on View being a real +# type rather than Any. +[[tool.mypy.overrides]] +module = ["lightly.data.sample"] +follow_imports = "normal" + +# tests/test_simclr_agrees.py imports examples/ and benchmarks/, which are not +# part of the wheel and are not passed to mypy on the command line. +[[tool.mypy.overrides]] +module = ["examples.*", "benchmarks.*"] +follow_imports = "skip" + [tool.pytest.ini_options] # importlib mode is required for the spawn-based DDP test pool (#1982): spawn # workers re-import the test module, which deadlocks under the default prepend diff --git a/tests/conftest.py b/tests/conftest.py index 009f0a9c6..dffc3de46 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,8 +1,16 @@ # content of conftest.py +import sys +from pathlib import Path + import pytest import torch.multiprocessing as mp +# tests/test_simclr_agrees.py imports examples/ and benchmarks/, neither of which +# is installed. --import-mode=importlib does not put the repo root on sys.path, +# so it goes here. +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + # Distributed (DDP) test pool, see #1982. The flag and gloo setup live in # tests/ddp_helpers.py so they can be typed and shared with the test modules; # the session hooks below only start/stop the pool. diff --git a/tests/test_simclr_agrees.py b/tests/test_simclr_agrees.py new file mode 100644 index 000000000..110f332e4 --- /dev/null +++ b/tests/test_simclr_agrees.py @@ -0,0 +1,252 @@ +"""The gate: examples/simclr.py and benchmarks/simclr/ are one method. + +The two files are written independently and restate each other on purpose. What +stops them drifting is not an import, it is this file. Every SimCLR in v1 failed +three of these: + + the two shipped heads differed by 4,922,112 parameters load_state_dict raised + split forward 2.582973, fused forward 2.757323 gradient cos sim 0.0837 + same seed, same image, 200 trials views differed on 108 + +What is compared and what is not +-------------------------------- +Compared: the block classes, the view contract, the fused forward, the head's +input width against the backbone's output, and that the temperature each file +declares is the one its criterion carries. + +Not compared: any value that belongs to the run rather than to the method. +Backbone family, image size, batch size, epochs, learning rate, the temperature's +value, the head's widths, blur, the normalization statistics and +``gather_distributed`` are all free to differ, and they do: the example is one +small configuration and the ``imagenet`` row is the paper's. + +Where a number has to be shared for a comparison to mean anything, the row is +built out of the example's own constants. That asserts the two files compute the +same thing given the same numbers. It never asserts that they chose the same +numbers. + +Not covered here: that each file excludes norms and biases from weight decay. +Both call ``lightly.optim.param_groups``, and ``tests/optim/test_param_groups.py`` +is where that rule is checked. +""" + +import random +from typing import Callable, List, NamedTuple + +import pytest +import torch +from PIL import Image +from torch import Tensor +from torch.nn import Linear, Module +from torch.testing import assert_close + +import examples.simclr as example +from benchmarks.simclr import benchmark +from benchmarks.simclr.datasets import SETTINGS, Setting +from lightly.backbones import TorchvisionResNetBackbone +from lightly.data.sample import Sample, View +from lightly.loss import NTXentLoss +from lightly.models.modules import SimCLRProjectionHead +from lightly.transforms import SimCLRTransform + +SEED = 11 + + +def setting_from_example() -> Setting: + """The row the example would be, if the example were a row. + + Deliberately not an entry in ``SETTINGS``: it exists so the two files can be + compared under one set of numbers, not so they share one. + """ + return Setting( + name="from_example", + size=example.INPUT_SIZE, + normalize=example.NORMALIZE, + stem="small_image", + blur=example.GAUSSIAN_BLUR, + color_jitter_strength=example.COLOR_JITTER_STRENGTH, + backbone="resnet18", + feature_dim=example.FEATURE_DIM, + hidden_dim=example.HIDDEN_DIM, + out_dim=example.OUT_DIM, + batch_size=example.BATCH_SIZE, + epochs=example.EPOCHS, + temperature=example.TEMPERATURE, + weight_decay=example.WEIGHT_DECAY, + optimizer="sgd", + lr=example.LR, + momentum=example.MOMENTUM, + warmup_epochs=0, + num_classes=10, + knn_k=200, + knn_t=0.1, + ) + + +def fixed_sample(seed: int, size: int, batch_size: int = 4) -> Sample: + torch.manual_seed(seed) + return Sample( + views=[View(torch.randn(batch_size, 3, size, size)) for _ in range(2)], + meta={"target": torch.zeros(batch_size, dtype=torch.long)}, + ) + + +def seed_everything(seed: int) -> None: + random.seed(seed) + torch.manual_seed(seed) + + +def synced(module: benchmark.SimCLR) -> benchmark.SimCLR: + module.backbone.load_state_dict(example.backbone.state_dict()) + module.head.load_state_dict(example.head.state_dict()) + return module + + +# --------------------------------------------------------------------------- # +# One method, given one set of numbers. +# --------------------------------------------------------------------------- # + + +def test_weights_are_interchangeable() -> None: + # v1: the two shipped heads differed by 4,922,112 parameters. + module = benchmark.SimCLR(setting_from_example()) + module.backbone.load_state_dict(example.backbone.state_dict()) + module.head.load_state_dict(example.head.state_dict()) + + +def test_same_loss_on_the_same_batch() -> None: + # v1: split forward 2.582973, fused forward 2.757323. + module = synced(benchmark.SimCLR(setting_from_example())) + sample = fixed_sample(SEED, size=example.INPUT_SIZE) + assert float(example.forward(sample).detach()) == pytest.approx( + float(module(sample).detach()), abs=1e-6 + ) + + +def test_same_views_from_the_same_seed() -> None: + # v1: the same seed and the same image gave different views on 108 of 200. + image = Image.new("RGB", (64, 64), color=(30, 90, 150)) + from_benchmark = benchmark.transform(setting_from_example()) + + seed_everything(SEED) + ours = example.transform(image) + seed_everything(SEED) + theirs = from_benchmark(image) + + assert len(ours) == len(theirs) + for a, b in zip(ours, theirs): + assert_close(a.data, b.data) + + +# --------------------------------------------------------------------------- # +# One method, each side reading its own numbers. +# --------------------------------------------------------------------------- # + + +class Side(NamedTuple): + name: str + backbone: TorchvisionResNetBackbone + head: Module + criterion: NTXentLoss + transform: SimCLRTransform + forward: Callable[[Sample], Tensor] + declared_temperature: float + size: int + + +def example_side() -> Side: + return Side( + name="example", + backbone=example.backbone, + head=example.head, + criterion=example.criterion, + transform=example.transform, + forward=example.forward, + declared_temperature=example.TEMPERATURE, + size=example.INPUT_SIZE, + ) + + +def benchmark_side() -> Side: + setting = SETTINGS["cifar10"] + module = benchmark.SimCLR(setting) + return Side( + name=f"benchmark[{setting.name}]", + backbone=module.backbone, + head=module.head, + criterion=module.criterion, + transform=benchmark.transform(setting), + forward=module.forward, + declared_temperature=setting.temperature, + size=setting.size, + ) + + +SIDES = [example_side, benchmark_side] +IDS = ["example", "benchmark"] + + +@pytest.mark.parametrize("build", SIDES, ids=IDS) +def test_the_same_blocks_are_assembled(build: Callable[[], Side]) -> None: + side = build() + assert isinstance(side.head, SimCLRProjectionHead) + assert isinstance(side.criterion, NTXentLoss) + assert isinstance(side.transform, SimCLRTransform) + assert hasattr(side.backbone, "embed") + + +@pytest.mark.parametrize("build", SIDES, ids=IDS) +def test_two_views_with_the_default_role(build: Callable[[], Side]) -> None: + side = build() + views = side.transform(Image.new("RGB", (64, 64))) + assert len(views) == 2 + assert [view.role for view in views] == ["view", "view"] + assert [view.stream for view in views] == ["image", "image"] + assert views[0].data.shape == views[1].data.shape + + +@pytest.mark.parametrize("build", SIDES, ids=IDS) +def test_the_head_reads_the_backbone_width(build: Callable[[], Side]) -> None: + side = build() + images = torch.randn(2, 3, side.size, side.size) + assert side.backbone.embed(images).shape[1] == side.backbone.feature_dim + first_linear = next( + module for module in side.head.modules() if isinstance(module, Linear) + ) + assert first_linear.in_features == side.backbone.feature_dim + + +@pytest.mark.parametrize("build", SIDES, ids=IDS) +def test_both_views_reach_the_encoder_in_one_call(build: Callable[[], Side]) -> None: + # BatchNorm has to see 2N. One call per view gives it N, and the gradients of + # the two agree at cosine similarity 0.0837. This drives each side's own + # forward, not encode: a test that calls encode itself proves nothing about + # the file it is meant to be checking. + side = build() + calls: List[int] = [] + handle = side.backbone.register_forward_pre_hook( + lambda _module, inputs: calls.append(inputs[0].size(0)) + ) + try: + side.forward(fixed_sample(SEED, size=side.size, batch_size=3)) + finally: + handle.remove() + assert calls == [6] + + +@pytest.mark.parametrize("build", SIDES, ids=IDS) +def test_the_temperature_is_stated_not_defaulted(build: Callable[[], Side]) -> None: + # v1: both examples took NTXentLoss()'s default 0.5 while the benchmark + # passed 0.1, and nothing noticed. + side = build() + assert side.criterion.temperature == side.declared_temperature + + +def test_every_row_states_a_width_its_head_can_read() -> None: + for name, setting in SETTINGS.items(): + module = benchmark.SimCLR(setting) + assert module.backbone.feature_dim == setting.feature_dim, name + first_linear = next( + child for child in module.head.modules() if isinstance(child, Linear) + ) + assert first_linear.in_features == setting.feature_dim, name