diffusion-insgen/training/contrastive_head.py

# Code are mainly borrowed from the official implementation of MoCo (https://github.com/facebookresearch/moco)

import numpy as np
import torch
import torch.nn as nn
from torch_utils import misc
from torch_utils import persistence

#----------------------------------------------------------------------------

# Contrastive head
@persistence.persistent_class
class CLHead(torch.nn.Module):
    def __init__(self,
        inplanes     = 256,   # Number of input features
        temperature  = 0.2,   # Temperature of logits
        queue_size   = 3500,  # Number of stored negative samples
        momentum     = 0.999, # Momentum for updating network
    ):
        super().__init__()
        self.inplanes = inplanes
        self.temperature = temperature
        self.queue_size = queue_size
        self.m = momentum

        self.mlp = nn.Sequential(nn.Linear(inplanes, inplanes), nn.ReLU(), nn.Linear(inplanes, 128))
        self.momentum_mlp = nn.Sequential(nn.Linear(inplanes, inplanes), nn.ReLU(), nn.Linear(inplanes, 128))
        self.momentum_mlp.requires_grad_(False)

        for param_q, param_k in zip(self.mlp.parameters(), self.momentum_mlp.parameters()):
            param_k.data.copy_(param_q.data)  # initialize
            param_k.requires_grad = False  # not update by gradient

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


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

    @torch.no_grad()
    def _dequeue_and_enqueue(self, keys):
        # gather keys before updating queue
        keys = concat_all_gather(keys)

        batch_size = keys.shape[0]
        keys = keys.T
        ptr = int(self.queue_ptr) 
        if batch_size > self.queue_size:
            self.queue[:, 0:] = keys[:, :self.queue_size]

        elif ptr + batch_size > self.queue_size:
            self.queue[:, ptr:] = keys[:, :self.queue_size - ptr]
            self.queue[:, :batch_size - (self.queue_size - ptr)] = keys[:, self.queue_size-ptr:]
            self.queue_ptr[0] = batch_size - (self.queue_size - ptr)
        else:
            self.queue[:, ptr:ptr + batch_size] = keys
            self.queue_ptr[0] = ptr + batch_size

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

        # If non-distributed now, return raw input directly.
        # We have no idea the effect of disabling shuffle BN to MoCo. 
        # Thus, we recommand train InsGen with more than 1 GPU always.
        if not torch.distributed.is_initialized():
            return x, torch.arange(x.shape[0])

        # gather from all gpus
        device = x.device
        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(device)

        # 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.
        """
        # If non-distributed now, return raw input directly.
        # We have no idea the effect of disabling shuffle BN to MoCo. 
        # Thus, we recommand train InsGen with more than 1 GPU always.
        if not torch.distributed.is_initialized():
            return x

        # 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(self, im_q, im_k, loss_only=False, update_q=False):
        """
        Input:
            im_q: a batch of query images
            im_k: a batch of key images
        Output:
            logits, targets
        """
        device = im_q.device
        im_q = im_q.to(torch.float32)
        im_k = im_k.to(torch.float32)
        # compute query features
        if im_q.ndim > 2:
            im_q = im_q.mean([2,3])
        q = self.mlp(im_q)  # 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
            if im_k.ndim > 2:
                im_k = im_k.mean([2,3])
            # shuffle for making use of BN
            im_k, idx_unshuffle = self._batch_shuffle_ddp(im_k)
            k = self.momentum_mlp(im_k)  # keys: NxC
            k = nn.functional.normalize(k, dim=1)

            # undo shuffle
            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()])

        # logits: Nx(1+K)
        logits = torch.cat([l_pos, l_neg], dim=1)

        # apply temperature
        logits /= self.temperature

        # labels: positive key indicators
        labels = torch.zeros(logits.shape[0], dtype=torch.long).cuda(device)

        # dequeue and enqueue
        if not loss_only:
            if update_q:
                with torch.no_grad():
                    temp_im_q, idx_unshuffle = self._batch_shuffle_ddp(im_q)
                    temp_q = self.momentum_mlp(temp_im_q)
                    temp_q = nn.functional.normalize(temp_q, dim=1)
                    temp_q = self._batch_unshuffle_ddp(temp_q, idx_unshuffle)
                    self._dequeue_and_enqueue(temp_q)
            else:
                self._dequeue_and_enqueue(k)

        # calculate loss
        loss = nn.functional.cross_entropy(logits, labels)

        return loss 

@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

#----------------------------------------------------------------------------