<gh_stars>0
from collections import OrderedDict

import torch
import torch.nn as nn
import torch.distributed as dist
from openselfsup.utils import print_log

from . import builder
from .registry import MODELS
from openselfsup.datasets import viz_utils
import random
import wandb

# TODO(cjrd) use this across all classes?

train_images_sent_to_wandb = False

class BaseModel(nn.Module):
    def __init__(self):
        super(BaseModel, self).__init__()
        self.debug = False

    def set_debug(self):
        self.debug = True

    def train_step(self, data, optimizer, **kwargs):
        losses = self(**data)
        loss, log_vars = self._parse_losses(losses)
        if 'img' in data:
            nimg = len(data['img'].data)
        else:
            nimg = len(data['img_q'].data)

        # Send a random image pair to wandb within this batch

        # The following is commented out as the data shape is the same as that the one are showing in wandb after the pipeline
        # global train_images_sent_to_wandb
        # try:
        #     # Here nimg represents a batch (default 64)
        #     if train_images_sent_to_wandb == False:
        #         rand_index = random.randint(0, nimg - 1)
        #         # The image data we get is a torch.Tensor
        #         # Type of data['img_q'].data[rand_index]:  <class 'torch.Tensor'> Shape is  torch.Size([10, 256, 256])
        #         # Type of data['img_k'].data[rand_index]:  <class 'torch.Tensor'> Shape is  torch.Size([2, 256, 256])
        #
        #         plt = viz_utils.read_msi_as_plt(data['img_k'].data[rand_index], data['img_q'].data[rand_index])
        #         wandb.log({'Images during train step': plt})
        #         train_images_sent_to_wandb = True
        #
        # except Exception as e: print(e)

        outputs = dict(loss=loss, log_vars=log_vars,
                       num_samples=nimg)
        return outputs

    def eval_step(self, data, optimizer, **kwargs):
        losses = self(**data)
        loss, log_vars = self._parse_losses(losses)
        if 'img' in data:
            nimg = len(data['img'].data)
        else:
            nimg = len(data['img_q'].data)
        outputs = dict(loss=loss, log_vars=log_vars,
                       num_samples=nimg)
        return outputs

    def val(self, data, optimizer, **kwargs):
        losses = self(**data)
        loss, log_vars = self._parse_losses(losses)
        if 'img' in data:
            nimg = len(data['img'].data)
        else:
            nimg = len(data['img_q'].data)
        outputs = dict(loss=loss, log_vars=log_vars,
                       num_samples=nimg)
        return outputs

    def _parse_losses(self, losses):
        """Parse the raw outputs (losses) of the network.
        Args:
            losses (dict): Raw output of the network, which usually contain
                losses and other necessary infomation.
        Returns:
            tuple[Tensor, dict]: (loss, log_vars), loss is the loss tensor \
                which may be a weighted sum of all losses, log_vars contains \
                all the variables to be sent to the logger.
        """
        log_vars = OrderedDict()
        for loss_name, loss_value in losses.items():
            if isinstance(loss_value, torch.Tensor):
                log_vars[loss_name] = loss_value.mean()
            elif isinstance(loss_value, list):
                log_vars[loss_name] = sum(_loss.mean() for _loss in loss_value)
            else:
                raise TypeError(
                    f'{loss_name} is not a tensor or list of tensors')

        loss = sum(_value for _key, _value in log_vars.items()
                   if 'loss' in _key)

        log_vars['loss'] = loss
        for loss_name, loss_value in log_vars.items():
            # reduce loss when distributed training
            if dist.is_available() and dist.is_initialized():
                loss_value = loss_value.data.clone()
                dist.all_reduce(loss_value.div_(dist.get_world_size()))
            log_vars[loss_name] = loss_value.item()

        return loss, log_vars


@MODELS.register_module
class MOCO(BaseModel):
    """MOCO.

    Implementation of "Momentum Contrast for Unsupervised Visual
    Representation Learning (https://arxiv.org/abs/1911.05722)".
    Part of the code is borrowed from:
    "https://github.com/facebookresearch/moco/blob/master/moco/builder.py".

    Args:
        backbone (dict): Config dict for module of backbone ConvNet.
        neck (dict): Config dict for module of deep features to compact feature vectors.
            Default: None.
        head (dict): Config dict for module of loss functions. Default: None.
        pretrained (str, optional): Path to pre-trained weights. Default: None.
        queue_len (int): Number of negative keys maintained in the queue.
            Default: 65536.
        feat_dim (int): Dimension of compact feature vectors. Default: 128.
        momentum (float): Momentum coefficient for the momentum-updated encoder.
            Default: 0.999.
    """

    def __init__(self,
                 backbone,
                 input_module_q=None,
                 input_module_k=None,
                 random_swap_q_k=False,
                 neck=None,
                 head=None,
                 pretrained=None,
                 queue_len=65536,
                 feat_dim=128,
                 momentum=0.999,
                 **kwargs):
        super(MOCO, self).__init__()

        # Allow for multi-band input
        self.input_module_q=builder.build_input_module(input_module_q)
        self.input_module_k=builder.build_input_module(input_module_k)
        self.random_swap_q_k = random_swap_q_k

        self.encoder_q = nn.Sequential(
            builder.build_backbone(backbone), builder.build_neck(neck))
        self.encoder_k = nn.Sequential(
            builder.build_backbone(backbone), builder.build_neck(neck))
        self.backbone = self.encoder_q[0]
        for param in self.encoder_k.parameters():
            param.requires_grad = False
        self.head = builder.build_head(head)

        self.init_weights(pretrained=pretrained)

        self.queue_len = queue_len
        self.momentum = momentum

        # create the queue
        self.register_buffer("queue", torch.randn(feat_dim, queue_len))
        self.queue = nn.functional.normalize(self.queue, dim=0)
        self.register_buffer("queue_ptr", torch.zeros(1, dtype=torch.long))

    def init_weights(self, pretrained=None):
        """Initialize the weights of model.

        Args:
            pretrained (str, optional): Path to pre-trained weights.
                Default: None.
        """
        if pretrained is not None:
            print_log('load model from: {}'.format(pretrained), logger='root')
        self.encoder_q[0].init_weights(pretrained=pretrained)
        self.encoder_q[1].init_weights(init_linear='kaiming')
        for param_q, param_k in zip(self.encoder_q.parameters(),
                                    self.encoder_k.parameters()):
            param_k.data.copy_(param_q.data)

    @torch.no_grad()
    def _momentum_update_key_encoder(self):
        """Momentum update of the key encoder."""
        for param_q, param_k in zip(self.encoder_q.parameters(),
                                    self.encoder_k.parameters()):
            param_k.data = param_k.data * self.momentum + \
                param_q.data * (1. - self.momentum)

    @torch.no_grad()
    def _dequeue_and_enqueue(self, keys):
        """Update queue."""
        # gather keys before updating queue
        if not self.debug:
            keys = concat_all_gather(keys)

        batch_size = keys.shape[0]

        ptr = int(self.queue_ptr)
        assert self.queue_len % batch_size == 0  # for simplicity

        # replace the keys at ptr (dequeue and enqueue)
        self.queue[:, ptr:ptr + batch_size] = keys.transpose(0, 1)
        ptr = (ptr + batch_size) % self.queue_len  # move pointer

        self.queue_ptr[0] = ptr

    @torch.no_grad()
    def _batch_shuffle_ddp(self, x):
        """Batch shuffle, for making use of BatchNorm.

        *** Only support DistributedDataParallel (DDP) model. ***
        """
        # gather from all gpus
        batch_size_this = x.shape[0]
        x_gather = concat_all_gather(x)
        batch_size_all = x_gather.shape[0]

        num_gpus = batch_size_all // batch_size_this

        # random shuffle index
        idx_shuffle = torch.randperm(batch_size_all).cuda()

        # broadcast to all gpus
        torch.distributed.broadcast(idx_shuffle, src=0)

        # index for restoring
        idx_unshuffle = torch.argsort(idx_shuffle)

        # shuffled index for this gpu
        gpu_idx = torch.distributed.get_rank()
        idx_this = idx_shuffle.view(num_gpus, -1)[gpu_idx]

        return x_gather[idx_this], idx_unshuffle

    @torch.no_grad()
    def _batch_unshuffle_ddp(self, x, idx_unshuffle):
        """Undo batch shuffle.

        *** Only support DistributedDataParallel (DDP) model. ***
        """
        # gather from all gpus
        batch_size_this = x.shape[0]
        x_gather = concat_all_gather(x)
        batch_size_all = x_gather.shape[0]

        num_gpus = batch_size_all // batch_size_this

        # restored index for this gpu
        gpu_idx = torch.distributed.get_rank()
        idx_this = idx_unshuffle.view(num_gpus, -1)[gpu_idx]

        return x_gather[idx_this]

    def forward_train(self, img=None, im_q=None, im_k=None, **kwargs):
        """Forward computation during training.

        Args:
            img (Tensor): Input of two concatenated images of shape (N, 2, C, H, W).
                Typically these should be mean centered and std scaled.

        Returns:
            dict[str, Tensor]: A dictionary of loss components.
        """
        if img is not None:
            assert img.dim() == 5, \
                "Input must have 5 dims, got: {}".format(img.dim())

        if img is not None and im_q is None:
            im_q = img[:, 0, ...].contiguous()
        if img is not None and im_k is None:
            im_k = img[:, 1, ...].contiguous()

        # compute query features
        if self.input_module_q is not None:
            im_q = self.input_module_q(im_q)
        if self.input_module_k is not None:
            im_k = self.input_module_k(im_k)

        # use images from both source in key and query
        half_bsz = im_q.shape[0]//2
        all_ims_qk = torch.cat([im_q[:half_bsz, ...], im_k[half_bsz:, ...]], dim=0)
        all_ims_kq = torch.cat([im_k[:half_bsz, ...], im_q[half_bsz:, ...]], dim=0)
        im_q = all_ims_qk
        im_k = all_ims_kq

        # randomly swap q and k  to avoid alignment bias
        if torch.rand(1).item() < 0.5:
            tmp = im_k
            im_k = im_q
            im_k = tmp

        q = self.encoder_q(im_q)[0]  # queries: NxC
        q = nn.functional.normalize(q, dim=1)

        # compute key features
        with torch.no_grad():  # no gradient to keys
            self._momentum_update_key_encoder()  # update the key encoder

            # shuffle for making use of BN
            if not self.debug:
                im_k, idx_unshuffle = self._batch_shuffle_ddp(im_k)

            k = self.encoder_k(im_k)[0]  # keys: NxC
            k = nn.functional.normalize(k, dim=1)

            # undo shuffle
            if not self.debug:
                k = self._batch_unshuffle_ddp(k, idx_unshuffle)

        # compute logits
        # Einstein sum is more intuitive
        # positive logits: Nx1
        l_pos = torch.einsum('nc,nc->n', [q, k]).unsqueeze(-1)
        # negative logits: NxK
        l_neg = torch.einsum('nc,ck->nk', [q, self.queue.clone().detach()])

        losses = self.head(l_pos, l_neg)
        self._dequeue_and_enqueue(k)

        return losses

    def forward_test(self, img, **kwargs):
        pass

    def forward(self, img=None, img_q=None, img_k=None, mode='train', **kwargs):
        """
        takes img (concat tensor of both images) or explicity img_q and img_k
        """
        if mode == 'train':
            return self.forward_train(img=img, im_q=img_q, im_k=img_k, **kwargs)
        elif mode == 'test':
            return self.forward_test(img=img, im_q=img_q, im_k=img_k, **kwargs)
        elif mode == 'extract':
            return self.backbone(img=img, im_q=img_q, im_k=img_k)
        else:
            raise Exception("No such mode: {}".format(mode))


# utils
@torch.no_grad()
def concat_all_gather(tensor):
    """Performs all_gather operation on the provided tensors.

    *** Warning ***: torch.distributed.all_gather has no gradient.
    """
    tensors_gather = [
        torch.ones_like(tensor)
        for _ in range(torch.distributed.get_world_size())
    ]
    torch.distributed.all_gather(tensors_gather, tensor, async_op=False)

    output = torch.cat(tensors_gather, dim=0)
    return output

from collections import OrderedDict

import torch
import torch.nn as nn
import torch.distributed as dist
from openselfsup.utils import print_log

from . import builder
from .registry import MODELS
from openselfsup.datasets import viz_utils
import random
import wandb

# TODO(cjrd) use this across all classes?

train_images_sent_to_wandb = False

class BaseModel(nn.Module):
    def __init__(self):
        super(BaseModel, self).__init__()
        self.debug = False

    def set_debug(self):
        self.debug = True

    def train_step(self, data, optimizer, **kwargs):
        losses = self(**data)
        loss, log_vars = self._parse_losses(losses)
        if 'img' in data:
            nimg = len(data['img'].data)
        else:
            nimg = len(data['img_q'].data)

        # Send a random image pair to wandb within this batch

        # The following is commented out as the data shape is the same as that the one are showing in wandb after the pipeline
        # global train_images_sent_to_wandb
        # try:
        #     # Here nimg represents a batch (default 64)
        #     if train_images_sent_to_wandb == False:
        #         rand_index = random.randint(0, nimg - 1)
        #         # The image data we get is a torch.Tensor
        #         # Type of data['img_q'].data[rand_index]:  <class 'torch.Tensor'> Shape is  torch.Size([10, 256, 256])
        #         # Type of data['img_k'].data[rand_index]:  <class 'torch.Tensor'> Shape is  torch.Size([2, 256, 256])
        #
        #         plt = viz_utils.read_msi_as_plt(data['img_k'].data[rand_index], data['img_q'].data[rand_index])
        #         wandb.log({'Images during train step': plt})
        #         train_images_sent_to_wandb = True
        #
        # except Exception as e: print(e)

        outputs = dict(loss=loss, log_vars=log_vars,
                       num_samples=nimg)
        return outputs

    def eval_step(self, data, optimizer, **kwargs):
        losses = self(**data)
        loss, log_vars = self._parse_losses(losses)
        if 'img' in data:
            nimg = len(data['img'].data)
        else:
            nimg = len(data['img_q'].data)
        outputs = dict(loss=loss, log_vars=log_vars,
                       num_samples=nimg)
        return outputs

    def val(self, data, optimizer, **kwargs):
        losses = self(**data)
        loss, log_vars = self._parse_losses(losses)
        if 'img' in data:
            nimg = len(data['img'].data)
        else:
            nimg = len(data['img_q'].data)
        outputs = dict(loss=loss, log_vars=log_vars,
                       num_samples=nimg)
        return outputs

    def _parse_losses(self, losses):
        """Parse the raw outputs (losses) of the network.
        Args:
            losses (dict): Raw output of the network, which usually contain
                losses and other necessary infomation.
        Returns:
            tuple[Tensor, dict]: (loss, log_vars), loss is the loss tensor \
                which may be a weighted sum of all losses, log_vars contains \
                all the variables to be sent to the logger.
        """
        log_vars = OrderedDict()
        for loss_name, loss_value in losses.items():
            if isinstance(loss_value, torch.Tensor):
                log_vars[loss_name] = loss_value.mean()
            elif isinstance(loss_value, list):
                log_vars[loss_name] = sum(_loss.mean() for _loss in loss_value)
            else:
                raise TypeError(
                    f'{loss_name} is not a tensor or list of tensors')

        loss = sum(_value for _key, _value in log_vars.items()
                   if 'loss' in _key)

        log_vars['loss'] = loss
        for loss_name, loss_value in log_vars.items():
            # reduce loss when distributed training
            if dist.is_available() and dist.is_initialized():
                loss_value = loss_value.data.clone()
                dist.all_reduce(loss_value.div_(dist.get_world_size()))
            log_vars[loss_name] = loss_value.item()

        return loss, log_vars


@MODELS.register_module
class MOCO(BaseModel):
    """MOCO.

    Implementation of "Momentum Contrast for Unsupervised Visual
    Representation Learning (https://arxiv.org/abs/1911.05722)".
    Part of the code is borrowed from:
    "https://github.com/facebookresearch/moco/blob/master/moco/builder.py".

    Args:
        backbone (dict): Config dict for module of backbone ConvNet.
        neck (dict): Config dict for module of deep features to compact feature vectors.
            Default: None.
        head (dict): Config dict for module of loss functions. Default: None.
        pretrained (str, optional): Path to pre-trained weights. Default: None.
        queue_len (int): Number of negative keys maintained in the queue.
            Default: 65536.
        feat_dim (int): Dimension of compact feature vectors. Default: 128.
        momentum (float): Momentum coefficient for the momentum-updated encoder.
            Default: 0.999.
    """

    def __init__(self,
                 backbone,
                 input_module_q=None,
                 input_module_k=None,
                 random_swap_q_k=False,
                 neck=None,
                 head=None,
                 pretrained=None,
                 queue_len=65536,
                 feat_dim=128,
                 momentum=0.999,
                 **kwargs):
        super(MOCO, self).__init__()

        # Allow for multi-band input
        self.input_module_q=builder.build_input_module(input_module_q)
        self.input_module_k=builder.build_input_module(input_module_k)
        self.random_swap_q_k = random_swap_q_k

        self.encoder_q = nn.Sequential(
            builder.build_backbone(backbone), builder.build_neck(neck))
        self.encoder_k = nn.Sequential(
            builder.build_backbone(backbone), builder.build_neck(neck))
        self.backbone = self.encoder_q[0]
        for param in self.encoder_k.parameters():
            param.requires_grad = False
        self.head = builder.build_head(head)

        self.init_weights(pretrained=pretrained)

        self.queue_len = queue_len
        self.momentum = momentum

        # create the queue
        self.register_buffer("queue", torch.randn(feat_dim, queue_len))
        self.queue = nn.functional.normalize(self.queue, dim=0)
        self.register_buffer("queue_ptr", torch.zeros(1, dtype=torch.long))

    def init_weights(self, pretrained=None):
        """Initialize the weights of model.

        Args:
            pretrained (str, optional): Path to pre-trained weights.
                Default: None.
        """
        if pretrained is not None:
            print_log('load model from: {}'.format(pretrained), logger='root')
        self.encoder_q[0].init_weights(pretrained=pretrained)
        self.encoder_q[1].init_weights(init_linear='kaiming')
        for param_q, param_k in zip(self.encoder_q.parameters(),
                                    self.encoder_k.parameters()):
            param_k.data.copy_(param_q.data)

    @torch.no_grad()
    def _momentum_update_key_encoder(self):
        """Momentum update of the key encoder."""
        for param_q, param_k in zip(self.encoder_q.parameters(),
                                    self.encoder_k.parameters()):
            param_k.data = param_k.data * self.momentum + \
                param_q.data * (1. - self.momentum)

    @torch.no_grad()
    def _dequeue_and_enqueue(self, keys):
        """Update queue."""
        # gather keys before updating queue
        if not self.debug:
            keys = concat_all_gather(keys)

        batch_size = keys.shape[0]

        ptr = int(self.queue_ptr)
        assert self.queue_len % batch_size == 0  # for simplicity

        # replace the keys at ptr (dequeue and enqueue)
        self.queue[:, ptr:ptr + batch_size] = keys.transpose(0, 1)
        ptr = (ptr + batch_size) % self.queue_len  # move pointer

        self.queue_ptr[0] = ptr

    @torch.no_grad()
    def _batch_shuffle_ddp(self, x):
        """Batch shuffle, for making use of BatchNorm.

        *** Only support DistributedDataParallel (DDP) model. ***
        """
        # gather from all gpus
        batch_size_this = x.shape[0]
        x_gather = concat_all_gather(x)
        batch_size_all = x_gather.shape[0]

        num_gpus = batch_size_all // batch_size_this

        # random shuffle index
        idx_shuffle = torch.randperm(batch_size_all).cuda()

        # broadcast to all gpus
        torch.distributed.broadcast(idx_shuffle, src=0)

        # index for restoring
        idx_unshuffle = torch.argsort(idx_shuffle)

        # shuffled index for this gpu
        gpu_idx = torch.distributed.get_rank()
        idx_this = idx_shuffle.view(num_gpus, -1)[gpu_idx]

        return x_gather[idx_this], idx_unshuffle

    @torch.no_grad()
    def _batch_unshuffle_ddp(self, x, idx_unshuffle):
        """Undo batch shuffle.

        *** Only support DistributedDataParallel (DDP) model. ***
        """
        # gather from all gpus
        batch_size_this = x.shape[0]
        x_gather = concat_all_gather(x)
        batch_size_all = x_gather.shape[0]

        num_gpus = batch_size_all // batch_size_this

        # restored index for this gpu
        gpu_idx = torch.distributed.get_rank()
        idx_this = idx_unshuffle.view(num_gpus, -1)[gpu_idx]

        return x_gather[idx_this]

    def forward_train(self, img=None, im_q=None, im_k=None, **kwargs):
        """Forward computation during training.

        Args:
            img (Tensor): Input of two concatenated images of shape (N, 2, C, H, W).
                Typically these should be mean centered and std scaled.

        Returns:
            dict[str, Tensor]: A dictionary of loss components.
        """
        if img is not None:
            assert img.dim() == 5, \
                "Input must have 5 dims, got: {}".format(img.dim())

        if img is not None and im_q is None:
            im_q = img[:, 0, ...].contiguous()
        if img is not None and im_k is None:
            im_k = img[:, 1, ...].contiguous()

        # compute query features
        if self.input_module_q is not None:
            im_q = self.input_module_q(im_q)
        if self.input_module_k is not None:
            im_k = self.input_module_k(im_k)

        # use images from both source in key and query
        half_bsz = im_q.shape[0]//2
        all_ims_qk = torch.cat([im_q[:half_bsz, ...], im_k[half_bsz:, ...]], dim=0)
        all_ims_kq = torch.cat([im_k[:half_bsz, ...], im_q[half_bsz:, ...]], dim=0)
        im_q = all_ims_qk
        im_k = all_ims_kq

        # randomly swap q and k  to avoid alignment bias
        if torch.rand(1).item() < 0.5:
            tmp = im_k
            im_k = im_q
            im_k = tmp

        q = self.encoder_q(im_q)[0]  # queries: NxC
        q = nn.functional.normalize(q, dim=1)

        # compute key features
        with torch.no_grad():  # no gradient to keys
            self._momentum_update_key_encoder()  # update the key encoder

            # shuffle for making use of BN
            if not self.debug:
                im_k, idx_unshuffle = self._batch_shuffle_ddp(im_k)

            k = self.encoder_k(im_k)[0]  # keys: NxC
            k = nn.functional.normalize(k, dim=1)

            # undo shuffle
            if not self.debug:
                k = self._batch_unshuffle_ddp(k, idx_unshuffle)

        # compute logits
        # Einstein sum is more intuitive
        # positive logits: Nx1
        l_pos = torch.einsum('nc,nc->n', [q, k]).unsqueeze(-1)
        # negative logits: NxK
        l_neg = torch.einsum('nc,ck->nk', [q, self.queue.clone().detach()])

        losses = self.head(l_pos, l_neg)
        self._dequeue_and_enqueue(k)

        return losses

    def forward_test(self, img, **kwargs):
        pass

    def forward(self, img=None, img_q=None, img_k=None, mode='train', **kwargs):
        """
        takes img (concat tensor of both images) or explicity img_q and img_k
        """
        if mode == 'train':
            return self.forward_train(img=img, im_q=img_q, im_k=img_k, **kwargs)
        elif mode == 'test':
            return self.forward_test(img=img, im_q=img_q, im_k=img_k, **kwargs)
        elif mode == 'extract':
            return self.backbone(img=img, im_q=img_q, im_k=img_k)
        else:
            raise Exception("No such mode: {}".format(mode))


# utils
@torch.no_grad()
def concat_all_gather(tensor):
    """Performs all_gather operation on the provided tensors.

    *** Warning ***: torch.distributed.all_gather has no gradient.
    """
    tensors_gather = [
        torch.ones_like(tensor)
        for _ in range(torch.distributed.get_world_size())
    ]
    torch.distributed.all_gather(tensors_gather, tensor, async_op=False)

    output = torch.cat(tensors_gather, dim=0)
    return output
# TODO(cjrd) use this across all classes? # Send a random image pair to wandb within this batch # The following is commented out as the data shape is the same as that the one are showing in wandb after the pipeline # global train_images_sent_to_wandb # try: #     # Here nimg represents a batch (default 64) #     if train_images_sent_to_wandb == False: #         rand_index = random.randint(0, nimg - 1) #         # The image data we get is a torch.Tensor #         # Type of data['img_q'].data[rand_index]:  <class 'torch.Tensor'> Shape is  torch.Size([10, 256, 256]) #         # Type of data['img_k'].data[rand_index]:  <class 'torch.Tensor'> Shape is  torch.Size([2, 256, 256]) # #         plt = viz_utils.read_msi_as_plt(data['img_k'].data[rand_index], data['img_q'].data[rand_index]) #         wandb.log({'Images during train step': plt}) #         train_images_sent_to_wandb = True # # except Exception as e: print(e) Parse the raw outputs (losses) of the network.         Args:             losses (dict): Raw output of the network, which usually contain                 losses and other necessary infomation.         Returns:             tuple[Tensor, dict]: (loss, log_vars), loss is the loss tensor \                 which may be a weighted sum of all losses, log_vars contains \                 all the variables to be sent to the logger. # reduce loss when distributed training MOCO.      Implementation of "Momentum Contrast for Unsupervised Visual     Representation Learning (https://arxiv.org/abs/1911.05722)".     Part of the code is borrowed from:     "https://github.com/facebookresearch/moco/blob/master/moco/builder.py".      Args:         backbone (dict): Config dict for module of backbone ConvNet.         neck (dict): Config dict for module of deep features to compact feature vectors.             Default: None.         head (dict): Config dict for module of loss functions. Default: None.         pretrained (str, optional): Path to pre-trained weights. Default: None.         queue_len (int): Number of negative keys maintained in the queue.             Default: 65536.         feat_dim (int): Dimension of compact feature vectors. Default: 128.         momentum (float): Momentum coefficient for the momentum-updated encoder.             Default: 0.999. # Allow for multi-band input # create the queue Initialize the weights of model.          Args:             pretrained (str, optional): Path to pre-trained weights.                 Default: None. Momentum update of the key encoder. Update queue. # gather keys before updating queue # for simplicity # replace the keys at ptr (dequeue and enqueue) # move pointer Batch shuffle, for making use of BatchNorm.          *** Only support DistributedDataParallel (DDP) model. *** # gather from all gpus # random shuffle index # broadcast to all gpus # index for restoring # shuffled index for this gpu Undo batch shuffle.          *** Only support DistributedDataParallel (DDP) model. *** # gather from all gpus # restored index for this gpu Forward computation during training.          Args:             img (Tensor): Input of two concatenated images of shape (N, 2, C, H, W).                 Typically these should be mean centered and std scaled.          Returns:             dict[str, Tensor]: A dictionary of loss components. # compute query features # use images from both source in key and query # randomly swap q and k  to avoid alignment bias # queries: NxC # compute key features # no gradient to keys # update the key encoder # shuffle for making use of BN # keys: NxC # undo shuffle # compute logits # Einstein sum is more intuitive # positive logits: Nx1 # negative logits: NxK takes img (concat tensor of both images) or explicity img_q and img_k # utils Performs all_gather operation on the provided tensors.      *** Warning ***: torch.distributed.all_gather has no gradient.