Repository: lightly-ai/lightly

Issue:
Error Using val_online_cls_top1
    print(f"max {metric}: {max(metric_callback.val_metrics[metric])}")
KeyError: 'val_online_cls_top1'


Discussion:
Hi, could you give some more detail on what you were trying to do?
Ok, I just used the `benchmarks/imagenet/resnet50/` example, with the following command: 
`python main.py --epochs 50 --num-classes 100 --train-dir ./data/train --val-dir ./data/val --methods swav --num-workers 8 --devices 1 --batch-size-per-device 64`
Then after pretraining this error popped up.
Did you modify the code in `swav.py`?

The error happens because it cannot find the `val_online_cls_top1` entry in the logs. `val_online_cls_top1` comes from the online classifier, did you change anything there? The value is logged here: https://github.com/lightly-ai/lightly/blob/c00f0669d987029032ff47b134f917bf5d97cae5/benchmarks/imagenet/resnet50/swav.py#L112-L115

and created here: https://github.com/lightly-ai/lightly/blob/c00f0669d987029032ff47b134f917bf5d97cae5/lightly/utils/benchmarking/online_linear_classifier.py#L45

For a temporary fix you can change the code here: https://github.com/lightly-ai/lightly/blob/c00f0669d987029032ff47b134f917bf5d97cae5/benchmarks/imagenet/resnet50/main.py#L218-L219

to this:
```
    for metric, values in sorted(metric_callback.val_metrics.items()):
        print(f"max {metric}: {max(values) if values else 0}")

```
I am now getting the following error: 

```  File "main.py", line 232, in <module>
    main(**vars(args))
  File "main.py", line 103, in main
    precision=precision,
  File "main.py", line 221, in pretrain
    val_dataloaders=val_dataloader,
  File "C:\Users\Lenovo\anaconda3\envs\lib\site-packages\pytorch_lightning\trainer\trainer.py", line 609, in fit
    self, self._fit_impl, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path
  File "C:\Users\Lenovo\anaconda3\envs\lib\site-packages\pytorch_lightning\trainer\call.py", line 38, in _call_and_handle_interrupt
    return trainer_fn(*args, **kwargs)
  File "C:\Users\Lenovo\anaconda3\envs\lib\site-packages\pytorch_lightning\trainer\trainer.py", line 650, in _fit_impl
    self._run(model, ckpt_path=self.ckpt_path)
  File "C:\Users\Lenovo\anaconda3\envs\lib\site-packages\pytorch_lightning\trainer\trainer.py", line 1112, in _run
    results = self._run_stage()
  File "C:\Users\Lenovo\anaconda3\envs\lib\site-packages\pytorch_lightning\trainer\trainer.py", line 1191, in _run_stage
    self._run_train()
  File "C:\Users\Lenovo\anaconda3\envs\lib\site-packages\pytorch_lightning\trainer\trainer.py", line 1214, in _run_train
    self.fit_loop.run()
  File "C:\Users\Lenovo\anaconda3\envs\lib\site-packages\pytorch_lightning\loops\loop.py", line 199, in run
    self.advance(*args, **kwargs)
  File "C:\Users\Lenovo\anaconda3\envs\lib\site-packages\pytorch_lightning\loops\fit_loop.py", line 267, in advance
    self._outputs = self.epoch_loop.run(self._data_fetcher)
  File "C:\Users\Lenovo\anaconda3\envs\lib\site-packages\pytorch_lightning\loops\loop.py", line 199, in run
    self.advance(*args, **kwargs)
  File "C:\Users\Lenovo\anaconda3\envs\lib\site-packages\pytorch_lightning\loops\epoch\training_epoch_loop.py", line 213, in advance
    batch_output = self.batch_loop.run(kwargs)
  File "C:\Users\Lenovo\anaconda3\envs\lib\site-packages\pytorch_lightning\loops\loop.py", line 199, in run
    self.advance(*args, **kwargs)
  File "C:\Users\Lenovo\anaconda3\envs\lib\site-packages\pytorch_lightning\loops\batch\training_batch_loop.py", line 88, in advance
    outputs = self.optimizer_loop.run(optimizers, kwargs)
  File "C:\Users\Lenovo\anaconda3\envs\lib\site-packages\pytorch_lightning\loops\loop.py", line 199, in run
    self.advance(*args, **kwargs)
  File "C:\Users\Lenovo\anaconda3\envs\lib\site-packages\pytorch_lightning\loops\optimization\optimizer_loop.py", line 202, in advance
    result = self._run_optimization(kwargs, self._optimizers[self.optim_progress.optimizer_position])
  File "C:\Users\Lenovo\anaconda3\envs\lib\site-packages\pytorch_lightning\loops\optimization\optimizer_loop.py", line 249, in _run_optimization
    self._optimizer_step(optimizer, opt_idx, kwargs.get("batch_idx", 0), closure)
  File "C:\Users\Lenovo\anaconda3\envs\lib\site-packages\pytorch_lightning\loops\optimization\optimizer_loop.py", line 379, in _optimizer_step
    using_lbfgs=is_lbfgs,
  File "C:\Users\Lenovo\anaconda3\envs\lib\site-packages\pytorch_lightning\trainer\trainer.py", line 1356, in _call_lightning_module_hook
    output = fn(*args, **kwargs)
  File "C:\Users\Lenovo\anaconda3\envs\lib\site-packages\pytorch_lightning\core\module.py", line 1754, in optimizer_step
    optimizer.step(closure=optimizer_closure)
  File "C:\Users\Lenovo\anaconda3\envs\lib\site-packages\pytorch_lightning\core\optimizer.py", line 169, in step
    step_output = self._strategy.optimizer_step(self._optimizer, self._optimizer_idx, closure, **kwargs)
  File "C:\Users\Lenovo\anaconda3\envs\lib\site-packages\pytorch_lightning\strategies\strategy.py", line 235, in optimizer_step
    optimizer, model=model, optimizer_idx=opt_idx, closure=closure, **kwargs
  File "C:\Users\Lenovo\anaconda3\envs\lib\site-packages\pytorch_lightning\plugins\precision\native_amp.py", line 75, in optimizer_step
    closure_result = closure()
  File "C:\Users\Lenovo\anaconda3\envs\lib\site-packages\pytorch_lightning\loops\optimization\optimizer_loop.py", line 149, in __call__
    self._result = self.closure(*args, **kwargs)
  File "C:\Users\Lenovo\anaconda3\envs\lib\site-packages\pytorch_lightning\loops\optimization\optimizer_loop.py", line 135, in closure
    step_output = self._step_fn()
  File "C:\Users\Lenovo\anaconda3\envs\lib\site-packages\pytorch_lightning\loops\optimization\optimizer_loop.py", line 419, in _training_step
    training_step_output = self.trainer._call_strategy_hook("training_step", *kwargs.values())
  File "C:\Users\Lenovo\anaconda3\envs\lib\site-packages\pytorch_lightning\trainer\trainer.py", line 1494, in _call_strategy_hook
    output = fn(*args, **kwargs)
  File "C:\Users\Lenovo\anaconda3\envs\lib\site-packages\pytorch_lightning\strategies\dp.py", line 134, in training_step
    return self.model(*args, **kwargs)
  File "C:\Users\Lenovo\anaconda3\envs\lib\site-packages\torch\nn\modules\module.py", line 1194, in _call_impl
    return forward_call(*input, **kwargs)
  File "C:\Users\Lenovo\anaconda3\envs\lib\site-packages\torch\nn\parallel\data_parallel.py", line 169, in forward
    return self.module(*inputs[0], **kwargs[0])
  File "C:\Users\Lenovo\anaconda3\envs\lib\site-packages\torch\nn\modules\module.py", line 1194, in _call_impl
    return forward_call(*input, **kwargs)
  File "C:\Users\Lenovo\anaconda3\envs\lib\site-packages\pytorch_lightning\overrides\data_parallel.py", line 77, in forward
    output = super().forward(*inputs, **kwargs)
  File "C:\Users\Lenovo\anaconda3\envs\lib\site-packages\pytorch_lightning\overrides\base.py", line 98, in forward
    output = self._forward_module.training_step(*inputs, **kwargs)
  File "H:\benchmarks_resnet50\swav.py", line 82, in training_step
    queues=self.queues,
  File "C:\Users\Lenovo\anaconda3\envs\lib\site-packages\torch\autograd\grad_mode.py", line 27, in decorate_context
    return func(*args, **kwargs)
  File "H:\benchmarks_resnet50\swav.py", line 195, in _update_queue
    _, projections = queues[i](projections[i], update=True)
  File "C:\Users\Lenovo\anaconda3\envs\lib\site-packages\torch\nn\modules\module.py", line 1194, in _call_impl
    return forward_call(*input, **kwargs)
  File "C:\Users\Lenovo\anaconda3\envs\lib\site-packages\lightly\loss\memory_bank.py", line 114, in forward
    _, dim = output.shape
ValueError: not enough values to unpack (expected 2, got 1)```
Could you show your code? I didn't have issues running the swav benchmark.
Sure, Here it is: 
```import math
from typing import List, Tuple

import torch
from pytorch_lightning import LightningModule
from torch import Tensor
from torch.nn import Identity, ModuleList
from torch.nn import functional as F
from torchvision.models import resnet50

from lightly.loss.memory_bank import MemoryBankModule
from lightly.loss.swav_loss import SwaVLoss
from lightly.models.modules import SwaVProjectionHead, SwaVPrototypes
from lightly.models.utils import get_weight_decay_parameters
from lightly.transforms import SwaVTransform
from lightly.utils.benchmarking import OnlineLinearClassifier
from lightly.utils.lars import LARS
from lightly.utils.scheduler import CosineWarmupScheduler

CROP_COUNTS: Tuple[int, int] = (2, 6)


class SwAV(LightningModule):
    def __init__(self, batch_size_per_device: int, num_classes: int) -> None:
        super().__init__()
        self.save_hyperparameters()
        self.batch_size_per_device = batch_size_per_device

        resnet = resnet50()
        resnet.fc = Identity()  # Ignore classification head
        self.backbone = resnet
        self.projection_head = SwaVProjectionHead()
        self.prototypes = SwaVPrototypes(n_steps_frozen_prototypes=1)
        self.criterion = SwaVLoss(sinkhorn_gather_distributed=True)
        self.online_classifier = OnlineLinearClassifier(num_classes=num_classes)

        # Use a queue for small batch sizes (<= 256).
        self.start_queue_at_epoch = 15
        self.n_batches_in_queue = 15
        self.queues = ModuleList(
            [
                MemoryBankModule(
                    size=self.n_batches_in_queue * self.batch_size_per_device
                )
                for _ in range(CROP_COUNTS[0])
            ]
        )

    def forward(self, x: Tensor) -> Tensor:
        return self.backbone(x)

    def project(self, x: Tensor) -> Tensor:
        x = self.projection_head(x)
        return F.normalize(x, dim=1, p=2)

    def training_step(
        self, batch: Tuple[List[Tensor], Tensor, List[str]], batch_idx: int
    ) -> Tensor:
        # Normalize the prototypes so they are on the unit sphere.
        self.prototypes.normalize()

        # The dataloader returns a list of image crops where the
        # first few items are high resolution crops and the rest are low
        # resolution crops.
        multi_crops, targets = batch[0], batch[1]

        # Forward pass through backbone and projection head.
        multi_crop_features = [
            self.forward(crops).flatten(start_dim=1) for crops in multi_crops
        ]
        multi_crop_projections = [
            self.project(features) for features in multi_crop_features
        ]

        # Get the queue projections and logits.
        queue_crop_logits = None
        with torch.no_grad():
            if self.current_epoch >= self.start_queue_at_epoch:
                # Start filling the queue.
                queue_crop_projections = _update_queue(
                    projections=multi_crop_projections[: CROP_COUNTS[0]],
                    queues=self.queues,
                )
                if batch_idx > self.n_batches_in_queue:
                    # The queue is filled, so we can start using it.
                    queue_crop_logits = [
                        self.prototypes(projections, step=self.current_epoch)
                        for projections in queue_crop_projections
                    ]

        # Get the rest of the multi-crop logits.
        multi_crop_logits = [
            self.prototypes(projections, step=self.current_epoch)
            for projections in multi_crop_projections
        ]

        # Calculate the SwAV loss.
        loss = self.criterion(
            high_resolution_outputs=multi_crop_logits[: CROP_COUNTS[0]],
            low_resolution_outputs=multi_crop_logits[CROP_COUNTS[0] :],
            queue_outputs=queue_crop_logits,
        )
        self.log(
            "train_loss",
            loss,
            prog_bar=True,
            sync_dist=True,)
           # batch_size_per_device=len(targets),
        #)

        # Calculate the classification loss.
        cls_loss, cls_log = self.online_classifier.training_step(
            (multi_crop_features[0].detach(), targets), batch_idx
        )
        #self.log_dict(cls_log, sync_dist=True, batch_size_per_device=len(targets)) ##rise error
        return loss + cls_loss

    def validation_step(
        self, batch: Tuple[Tensor, Tensor, List[str]], batch_idx: int
    ) -> Tensor:
        images, targets = batch[0], batch[1]
        features = self.forward(images).flatten(start_dim=1)
        cls_loss, cls_log = self.online_classifier.validation_step(
            (features.detach(), targets), batch_idx
        )
#         self.log_dict(
#             cls_log, prog_bar=True, sync_dist=True, batch_size_per_device=len(targets)
#         ) ##rise error
        return cls_loss

    def configure_optimizers(self):
        # Don't use weight decay for batch norm, bias parameters, and classification
        # head to improve performance.
        params, params_no_weight_decay = get_weight_decay_parameters(
            [self.backbone, self.projection_head, self.prototypes]
        )
        optimizer = LARS(
            [
                {"name": "swav", "params": params},
                {
                    "name": "swav_no_weight_decay",
                    "params": params_no_weight_decay,
                    "weight_decay": 0.0,
                },
                {
                    "name": "online_classifier",
                    "params": self.online_classifier.parameters(),
                    "weight_decay": 0.0,
                },
            ],
            # Smaller learning rate for smaller batches: lr=0.6 for batch_size=256
            # scaled linearly by batch size to lr=4.8 for batch_size=2048.
            # See Appendix A.1. and A.6. in SwAV paper https://arxiv.org/pdf/2006.09882.pdf
            lr=0.6 * (self.batch_size_per_device * self.trainer.world_size) / 256,
            momentum=0.9,
            weight_decay=1e-6,
        )
        scheduler = {
            "scheduler": CosineWarmupScheduler(
                optimizer=optimizer,
                warmup_epochs=(
                    self.trainer.estimated_stepping_batches
                    / self.trainer.max_epochs
                    * 10
                ),
                max_epochs=self.trainer.estimated_stepping_batches,
                end_value=0.0006
                * (self.batch_size_per_device * self.trainer.world_size)
                / 256,
            ),
            "interval": "step",
        }
        return [optimizer], [scheduler]


transform = SwaVTransform(crop_counts=CROP_COUNTS)


@torch.no_grad()
def _update_queue(
    projections: List[Tensor],
    queues: ModuleList,
):
    """Adds the high resolution projections to the queues and returns the queues."""

    if len(projections) != len(queues):
        raise ValueError(
            f"The number of queues ({len(queues)}) should be equal to the number of high "
            f"resolution inputs ({len(projections)})."
        )

    # Get the queue projections
    queue_projections = []
    for i in range(len(queues)):
        _, projections = queues[i](projections[i], update=True)
        # Queue projections are in (num_ftrs X queue_length) shape, while the high res
        # projections are in (batch_size_per_device X num_ftrs). Swap the axes for interoperability.
        projections = torch.permute(projections, (1, 0))
        queue_projections.append(projections)

    return queue_projections```


Resolving patch:
diff --git a/benchmarks/imagenet/resnet50/swav.py b/benchmarks/imagenet/resnet50/swav.py
--- a/benchmarks/imagenet/resnet50/swav.py
+++ b/benchmarks/imagenet/resnet50/swav.py
@@ -105,14 +105,14 @@ def training_step(
             loss,
             prog_bar=True,
             sync_dist=True,
-            batch_size_per_device=len(targets),
+            batch_size=len(targets),
         )
 
         # Calculate the classification loss.
         cls_loss, cls_log = self.online_classifier.training_step(
             (multi_crop_features[0].detach(), targets), batch_idx
         )
-        self.log_dict(cls_log, sync_dist=True, batch_size_per_device=len(targets))
+        self.log_dict(cls_log, sync_dist=True, batch_size=len(targets))
         return loss + cls_loss
 
     def validation_step(
@@ -123,9 +123,7 @@ def validation_step(
         cls_loss, cls_log = self.online_classifier.validation_step(
             (features.detach(), targets), batch_idx
         )
-        self.log_dict(
-            cls_log, prog_bar=True, sync_dist=True, batch_size_per_device=len(targets)
-        )
+        self.log_dict(cls_log, prog_bar=True, sync_dist=True, batch_size=len(targets))
         return cls_loss
 
     def configure_optimizers(self):
@@ -192,10 +190,10 @@ def _update_queue(
     # Get the queue projections
     queue_projections = []
     for i in range(len(queues)):
-        _, projections = queues[i](projections[i], update=True)
+        _, queue_proj = queues[i](projections[i], update=True)
         # Queue projections are in (num_ftrs X queue_length) shape, while the high res
         # projections are in (batch_size_per_device X num_ftrs). Swap the axes for interoperability.
-        projections = torch.permute(projections, (1, 0))
-        queue_projections.append(projections)
+        queue_proj = torch.permute(queue_proj, (1, 0))
+        queue_projections.append(queue_proj)
 
     return queue_projections


Test patch: