Source code for mmselfsup.models.algorithms.swav
# Copyright (c) OpenMMLab. All rights reserved.
from typing import Dict, List, Tuple
import torch
from mmselfsup.registry import MODELS
from mmselfsup.structures import SelfSupDataSample
from .base import BaseModel
[docs]@MODELS.register_module()
class SwAV(BaseModel):
"""SwAV.
Implementation of `Unsupervised Learning of Visual Features by Contrasting
Cluster Assignments <https://arxiv.org/abs/2006.09882>`_. The queue is
built in `engine/hooks/swav_hook.py`.
"""
[docs] def extract_feat(self, inputs: List[torch.Tensor],
**kwargs) -> Tuple[torch.Tensor]:
"""Function to extract features from backbone.
Args:
inputs (List[torch.Tensor]): The input images.
Returns:
Tuple[torch.Tensor]: backbone outputs.
"""
x = self.backbone(inputs[0])
return x
[docs] def loss(self, inputs: List[torch.Tensor],
data_samples: List[SelfSupDataSample],
**kwargs) -> Dict[str, torch.Tensor]:
"""Forward computation during training.
Args:
inputs (List[torch.Tensor]): The input images.
data_samples (List[SelfSupDataSample]): All elements required
during the forward function.
Returns:
Dict[str, torch.Tensor]: A dictionary of loss components.
"""
assert isinstance(inputs, list)
# multi-res forward passes
idx_crops = torch.cumsum(
torch.unique_consecutive(
torch.tensor([input.shape[-1] for input in inputs]),
return_counts=True)[1], 0)
start_idx = 0
output = []
for end_idx in idx_crops:
_out = self.backbone(torch.cat(inputs[start_idx:end_idx]))
output.append(_out)
start_idx = end_idx
output = self.neck(output)[0]
loss = self.head(output)
losses = dict(loss=loss)
return losses