Repository: lightly-ai/lightly

Issue:
Add callback for knn-validation
Add hooks for validation at the end of each epoch (only informative if there are labels). Let's make sure it can be switched off and that the user can pass a validation set of his choice.


Discussion:
I can handle this after I finish #61, or potentially before if we end up using this to reproduce some of the results (CIFAR-10) in the SimSiam paper
Let's go one-by-one because we already have an internal version of this with which we can evaluate whether SimSiam works as expected.
I don't like callbacks for that. It would just blow up the modules. Not everyone wants to run kNN-validation and needs this. You could only use it for image classification and only if you have labels. 
Why not adding some utility methods instead to quickly build one yourself and provide the sample code? That way it could be easier to adapt it.
I would agree we shouldn't dump every little thing we think up into its own module, however I do think offering some basic callbacks would not be inappropriate. For example, PL has a callback for an [online evaluator](https://pytorch-lightning-bolts.readthedocs.io/en/stable/self_supervised_callbacks.html). 

Resolving patch:
diff --git a/docs/source/getting_started/benchmarks/cifar10_benchmark.py b/docs/source/getting_started/benchmarks/cifar10_benchmark.py
--- a/docs/source/getting_started/benchmarks/cifar10_benchmark.py
+++ b/docs/source/getting_started/benchmarks/cifar10_benchmark.py
@@ -1,6 +1,13 @@
 # -*- coding: utf-8 -*-
 """
 
+Note that this benchmark also supports a multi-GPU setup. If you run it on
+a system with multiple GPUs make sure that you kill all the processes when
+killing the application. Due to the way we setup this benchmark the distributed
+processes might continue the benchmark if one of the nodes is killed.
+If you know how to fix this don't hesitate to create an issue or PR :)
+
+
 Code to reproduce the benchmark results:
 
 | Model   | Epochs | Batch Size | Test Accuracy | Peak GPU usage |
@@ -23,6 +30,7 @@
 import numpy as np
 import pytorch_lightning as pl
 import lightly
+from lightly.utils import BenchmarkModule
 
 num_workers = 8
 memory_bank_size = 4096
@@ -38,7 +46,8 @@
 batch_sizes = [128, 512]
 
 # use a GPU if available
-gpus = 1 if torch.cuda.is_available() else 0
+gpus = -1 if torch.cuda.is_available() else 0
+distributed_backend = 'ddp' if torch.cuda.device_count() > 1 else None
 
 # Adapted from our MoCo Tutorial on CIFAR-10
 #
@@ -128,107 +137,11 @@ def get_data_loaders(batch_size: int):
     )
 
     return dataloader_train_ssl, dataloader_train_kNN, dataloader_test
-
-# code for kNN prediction from here:
-# https://colab.research.google.com/github/facebookresearch/moco/blob/colab-notebook/colab/moco_cifar10_demo.ipynb
-def knn_predict(feature, feature_bank, feature_labels, classes: int, knn_k: int, knn_t: float):
-    """Helper method to run kNN predictions on features based on a feature bank
-
-    Args:
-        feature: Tensor of shape [N, D] consisting of N D-dimensional features
-        feature_bank: Tensor of a database of features used for kNN
-        feature_labels: Labels for the features in our feature_bank
-        classes: Number of classes (e.g. 10 for CIFAR-10)
-        knn_k: Number of k neighbors used for kNN
-        knn_t: 
-
-    """
-
-    # compute cos similarity between each feature vector and feature bank ---> [B, N]
-    sim_matrix = torch.mm(feature, feature_bank)
-    # [B, K]
-    sim_weight, sim_indices = sim_matrix.topk(k=knn_k, dim=-1)
-    # [B, K]
-    sim_labels = torch.gather(feature_labels.expand(feature.size(0), -1), dim=-1, index=sim_indices)
-
-    # we do a reweighting of the similarities 
-    sim_weight = (sim_weight / knn_t).exp()
-
-    # counts for each class
-    one_hot_label = torch.zeros(feature.size(0) * knn_k, classes, device=sim_labels.device)
-    # [B*K, C]
-    one_hot_label = one_hot_label.scatter(dim=-1, index=sim_labels.view(-1, 1), value=1.0)
-    # weighted score ---> [B, C]
-    pred_scores = torch.sum(one_hot_label.view(feature.size(0), -1, classes) * sim_weight.unsqueeze(dim=-1), dim=1)
-
-    pred_labels = pred_scores.argsort(dim=-1, descending=True)
-    return pred_labels
-
-
-class BenchmarkModule(pl.LightningModule):
-    """A PyTorch Lightning Module for automated kNN callback
-    
-    At the end of every training epoch we create a feature bank by inferencing
-    the backbone on the dataloader passed to the module. 
-    At every validation step we predict features on the validation data.
-    After all predictions on validation data (validation_epoch_end) we evaluate
-    the predictions on a kNN classifier on the validation data using the 
-    feature_bank features from the train data.
-    We can access the highest accuracy during a kNN prediction using the 
-    max_accuracy attribute.
-    """
-    def __init__(self, dataloader_kNN):
-        super().__init__()
-        self.backbone = nn.Module()
-        self.max_accuracy = 0.0
-        self.dataloader_kNN = dataloader_kNN
-
-    def training_epoch_end(self, outputs):
-        # update feature bank at the end of each training epoch
-        self.backbone.eval()
-        self.feature_bank = []
-        self.targets_bank = []
-        with torch.no_grad():
-            for data in self.dataloader_kNN:
-                img, target, _ = data
-                if gpus > 0:
-                    img = img.cuda()
-                    target = target.cuda()
-                feature = self.backbone(img).squeeze()
-                feature = F.normalize(feature, dim=1)
-                self.feature_bank.append(feature)
-                self.targets_bank.append(target)
-        self.feature_bank = torch.cat(self.feature_bank, dim=0).t().contiguous()
-        self.targets_bank = torch.cat(self.targets_bank, dim=0).t().contiguous()
-        self.backbone.train()
-
-    def validation_step(self, batch, batch_idx):
-        # we can only do kNN predictions once we have a feature bank
-        if hasattr(self, 'feature_bank') and hasattr(self, 'targets_bank'):
-            images, targets, _ = batch
-            feature = self.backbone(images).squeeze()
-            feature = F.normalize(feature, dim=1)
-            pred_labels = knn_predict(feature, self.feature_bank, self.targets_bank, classes, knn_k, knn_t)
-            num = images.size(0)
-            top1 = (pred_labels[:, 0] == targets).float().sum().item()
-            return (num, top1)
     
-    def validation_epoch_end(self, outputs):
-        if outputs:
-            total_num = 0
-            total_top1 = 0.
-            for (num, top1) in outputs:
-                total_num += num
-                total_top1 += top1
-            acc = float(total_top1 / total_num)
-            if acc > self.max_accuracy:
-                self.max_accuracy = acc
-            self.log('kNN_accuracy', acc * 100.0)
-
 
 class MocoModel(BenchmarkModule):
-    def __init__(self, dataloader_kNN):
-        super().__init__(dataloader_kNN)
+    def __init__(self, dataloader_kNN, num_classes):
+        super().__init__(dataloader_kNN, num_classes)
         # create a ResNet backbone and remove the classification head
         resnet = lightly.models.ResNetGenerator('resnet-18', num_splits=8)
         self.backbone = nn.Sequential(
@@ -264,8 +177,8 @@ def configure_optimizers(self):
 
 
 class SimCLRModel(BenchmarkModule):
-    def __init__(self, dataloader_kNN):
-        super().__init__(dataloader_kNN)
+    def __init__(self, dataloader_kNN, num_classes):
+        super().__init__(dataloader_kNN, num_classes)
         # create a ResNet backbone and remove the classification head
         resnet = lightly.models.ResNetGenerator('resnet-18')
         self.backbone = nn.Sequential(
@@ -295,8 +208,8 @@ def configure_optimizers(self):
 
 
 class SimSiamModel(BenchmarkModule):
-    def __init__(self, dataloader_kNN):
-        super().__init__(dataloader_kNN)
+    def __init__(self, dataloader_kNN, num_classes):
+        super().__init__(dataloader_kNN, num_classes)
         # create a ResNet backbone and remove the classification head
         resnet = lightly.models.ResNetGenerator('resnet-18')
         self.backbone = nn.Sequential(
@@ -333,25 +246,26 @@ def configure_optimizers(self):
 
 # loop through configurations and train models
 for batch_size in batch_sizes:
-    for Model in models:
+    for BenchmarkModel in models:
         runs = []
         for seed in range(n_runs):
             pl.seed_everything(seed)
             dataloader_train_ssl, dataloader_train_kNN, dataloader_test = get_data_loaders(batch_size)
-            model = Model(dataloader_train_kNN)
+            benchmark_model = BenchmarkModel(dataloader_train_kNN, classes)
             trainer = pl.Trainer(max_epochs=max_epochs, gpus=gpus,
-                                progress_bar_refresh_rate=100)
+                                progress_bar_refresh_rate=100,
+                                distributed_backend=distributed_backend)
             trainer.fit(
-                model,
+                benchmark_model,
                 train_dataloader=dataloader_train_ssl,
                 val_dataloaders=dataloader_test
             )
             gpu_memory_usage.append(torch.cuda.max_memory_allocated())
             torch.cuda.reset_peak_memory_stats()
-            runs.append(model.max_accuracy)
+            runs.append(benchmark_model.max_accuracy)
 
             # delete model and trainer + free up cuda memory
-            del model
+            del benchmark_model
             del trainer
             torch.cuda.empty_cache()
         bench_results.append(runs)
diff --git a/lightly/utils/__init__.py b/lightly/utils/__init__.py
--- a/lightly/utils/__init__.py
+++ b/lightly/utils/__init__.py
@@ -13,3 +13,5 @@
 from lightly.utils.io import load_embeddings
 from lightly.utils.io import load_embeddings_as_dict
 from lightly.utils.embeddings_2d import fit_pca
+from lightly.utils.benchmarking import BenchmarkModule
+from lightly.utils.benchmarking import knn_predict
diff --git a/lightly/utils/benchmarking.py b/lightly/utils/benchmarking.py
new file mode 100644
--- /dev/null
+++ b/lightly/utils/benchmarking.py
@@ -0,0 +1,221 @@
+""" Helper modules for benchmarking SSL models """
+
+# Copyright (c) 2020. Lightly AG and its affiliates.
+# All Rights Reserved
+
+import torch
+import torch.nn as nn
+import torch.distributed as dist 
+from torch.utils.data import DataLoader
+import torch.nn.functional as F
+import pytorch_lightning as pl
+
+# code for kNN prediction from here:
+# https://colab.research.google.com/github/facebookresearch/moco/blob/colab-notebook/colab/moco_cifar10_demo.ipynb
+
+
+def knn_predict(feature: torch.Tensor,
+                feature_bank: torch.Tensor,
+                feature_labels: torch.Tensor, 
+                num_classes: int,
+                knn_k: int=200,
+                knn_t: float=0.1) -> torch.Tensor:
+    """Run kNN predictions on features based on a feature bank
+
+    This method is commonly used to monitor performance of self-supervised
+    learning methods.
+
+    The default parameters are the ones
+    used in https://arxiv.org/pdf/1805.01978v1.pdf.
+
+    Args:
+        feature: 
+            Tensor of shape [N, D] for which you want predictions
+        feature_bank: 
+            Tensor of a database of features used for kNN
+        feature_labels: 
+            Labels for the features in our feature_bank
+        num_classes: 
+            Number of classes (e.g. `10` for CIFAR-10)
+        knn_k: 
+            Number of k neighbors used for kNN
+        knn_t: 
+            Temperature parameter to reweights similarities for kNN
+
+    Returns:
+        A tensor containing the kNN predictions
+
+    Examples:
+        >>> images, targets, _ = batch
+        >>> feature = backbone(images).squeeze()
+        >>> # we recommend to normalize the features
+        >>> feature = F.normalize(feature, dim=1)
+        >>> pred_labels = knn_predict(
+        >>>     feature,
+        >>>     feature_bank,
+        >>>     targets_bank,
+        >>>     num_classes=10,
+        >>> )
+    """
+
+    # compute cos similarity between each feature vector and feature bank ---> [B, N]
+    sim_matrix = torch.mm(feature, feature_bank)
+    # [B, K]
+    sim_weight, sim_indices = sim_matrix.topk(k=knn_k, dim=-1)
+    # [B, K]
+    sim_labels = torch.gather(feature_labels.expand(
+        feature.size(0), -1), dim=-1, index=sim_indices)
+    # we do a reweighting of the similarities
+    sim_weight = (sim_weight / knn_t).exp()
+    # counts for each class
+    one_hot_label = torch.zeros(feature.size(
+        0) * knn_k, num_classes, device=sim_labels.device)
+    # [B*K, C]
+    one_hot_label = one_hot_label.scatter(
+        dim=-1, index=sim_labels.view(-1, 1), value=1.0)
+    # weighted score ---> [B, C]
+    pred_scores = torch.sum(one_hot_label.view(feature.size(
+        0), -1, num_classes) * sim_weight.unsqueeze(dim=-1), dim=1)
+    pred_labels = pred_scores.argsort(dim=-1, descending=True)
+    return pred_labels
+
+
+class BenchmarkModule(pl.LightningModule):
+    """A PyTorch Lightning Module for automated kNN callback
+
+    At the end of every training epoch we create a feature bank by feeding the
+    `dataloader_kNN` passed to the module through the backbone.
+    At every validation step we predict features on the validation data.
+    After all predictions on validation data (validation_epoch_end) we evaluate
+    the predictions on a kNN classifier on the validation data using the
+    feature_bank features from the train data.
+
+    We can access the highest test accuracy during a kNN prediction 
+    using the `max_accuracy` attribute.
+
+    Attributes:
+        backbone:
+            The backbone model used for kNN validation. Make sure that you set the
+            backbone when inheriting from `BenchmarkModule`.
+        max_accuracy:
+            Floating point number between 0.0 and 1.0 representing the maximum
+            test accuracy the benchmarked model has achieved.
+        dataloader_kNN:
+            Dataloader to be used after each training epoch to create feature bank.
+        num_classes:
+            Number of classes. E.g. for cifar10 we have 10 classes. (default: 10)
+        knn_k:
+            Number of nearest neighbors for kNN
+        knn_t:
+            Temperature parameter for kNN
+
+    Examples:
+        >>> class SimSiamModel(BenchmarkingModule):
+        >>>     def __init__(dataloader_kNN, num_classes):
+        >>>         super().__init__(dataloader_kNN, num_classes)
+        >>>         resnet = lightly.models.ResNetGenerator('resnet-18')
+        >>>         self.backbone = nn.Sequential(
+        >>>             *list(resnet.children())[:-1],
+        >>>             nn.AdaptiveAvgPool2d(1),
+        >>>         )
+        >>>         self.resnet_simsiam = 
+        >>>             lightly.models.SimSiam(self.backbone, num_ftrs=512)
+        >>>         self.criterion = lightly.loss.SymNegCosineSimilarityLoss()
+        >>>
+        >>>     def forward(self, x):
+        >>>         self.resnet_simsiam(x)
+        >>>
+        >>>     def training_step(self, batch, batch_idx):
+        >>>         (x0, x1), _, _ = batch
+        >>>         x0, x1 = self.resnet_simsiam(x0, x1)
+        >>>         loss = self.criterion(x0, x1)
+        >>>         return loss
+        >>>     def configure_optimizers(self):
+        >>>         optim = torch.optim.SGD(
+        >>>             self.resnet_simsiam.parameters(), lr=6e-2, momentum=0.9
+        >>>         )
+        >>>         return [optim]
+        >>>
+        >>> model = SimSiamModel(dataloader_train_kNN)
+        >>> trainer = pl.Trainer()
+        >>> trainer.fit(
+        >>>     model,
+        >>>     train_dataloader=dataloader_train_ssl,
+        >>>     val_dataloaders=dataloader_test
+        >>> )
+        >>> # you can get the peak accuracy using
+        >>> print(model.max_accuracy)
+
+    """
+
+    def __init__(self,
+                 dataloader_kNN: DataLoader,
+                 num_classes: int,
+                 knn_k: int=200,
+                 knn_t: float=0.1):
+        super().__init__()
+        self.backbone = nn.Module()
+        self.max_accuracy = 0.0
+        self.dataloader_kNN = dataloader_kNN
+        self.num_classes = num_classes
+        self.knn_k = knn_k
+        self.knn_t = knn_t
+
+        # create dummy param to keep track of the device the model is using
+        self.dummy_param = nn.Parameter(torch.empty(0))
+
+    def training_epoch_end(self, outputs):
+        # update feature bank at the end of each training epoch
+        self.backbone.eval()
+        self.feature_bank = []
+        self.targets_bank = []
+        with torch.no_grad():
+            for data in self.dataloader_kNN:
+                img, target, _ = data
+                img = img.to(self.dummy_param.device)
+                target = target.to(self.dummy_param.device)
+                feature = self.backbone(img).squeeze()
+                feature = F.normalize(feature, dim=1)
+                self.feature_bank.append(feature)
+                self.targets_bank.append(target)
+        self.feature_bank = torch.cat(
+            self.feature_bank, dim=0).t().contiguous()
+        self.targets_bank = torch.cat(
+            self.targets_bank, dim=0).t().contiguous()
+        self.backbone.train()
+
+    def validation_step(self, batch, batch_idx):
+        # we can only do kNN predictions once we have a feature bank
+        if hasattr(self, 'feature_bank') and hasattr(self, 'targets_bank'):
+            images, targets, _ = batch
+            feature = self.backbone(images).squeeze()
+            feature = F.normalize(feature, dim=1)
+            pred_labels = knn_predict(
+                feature,
+                self.feature_bank,
+                self.targets_bank,
+                self.num_classes,
+                self.knn_k,
+                self.knn_t
+            )
+            num = images.size()
+            top1 = (pred_labels[:, 0] == targets).float().sum()
+            return (num, top1)
+
+    def validation_epoch_end(self, outputs):
+        device = self.dummy_param.device
+        if outputs:
+            total_num = torch.Tensor([0]).to(device)
+            total_top1 = torch.Tensor([0.]).to(device)
+            for (num, top1) in outputs:
+                total_num += num[0]
+                total_top1 += top1
+
+            if dist.get_world_size() > 1:
+                dist.all_reduce(total_num)
+                dist.all_reduce(total_top1)
+
+            acc = float(total_top1.item() / total_num.item())
+            if acc > self.max_accuracy:
+                self.max_accuracy = acc
+            self.log('kNN_accuracy', acc * 100.0, prog_bar=True)


Test patch: