xdecoder/Instruct-X-Decoder

1import torch2import torch.nn as nn3import torch.nn.functional as F4 5def is_dist_initialized():6    return torch.distributed.is_initialized()7 8def get_world_size():9    if is_dist_initialized():10        return torch.distributed.get_world_size()11    return 112 13def all_gather_grad(x):14    if get_world_size() > 1:15        all_x = [torch.zeros_like(x) for _ in range(get_world_size())]16        torch.distributed.all_gather(all_x, x)17        all_x[torch.distributed.get_rank()] = x18        x = torch.cat(all_x, dim=0)19    return x20 21@torch.no_grad()22def all_gather_nograd(tensor):23    # from albef24    """25    Performs all_gather operation on the provided tensors.26    *** Warning ***: torch.distributed.all_gather has no gradient.27    """28    if get_world_size() > 1:29        tensors_gather = [torch.ones_like(tensor)30            for _ in range(torch.distributed.get_world_size())]31        torch.distributed.all_gather(tensors_gather, tensor, async_op=False)32 33        tensor = torch.cat(tensors_gather, dim=0)34    return tensor35 36def image_text_contrastive_loss(image_feat, text_feat, temperature, image_id=None, text_id=None):37    # add the following 4 lines38    image_feat = all_gather_grad(image_feat)39    text_feat = all_gather_grad(text_feat)40    41    logits = torch.matmul(image_feat, text_feat.t())42    logits /= temperature43    44    if image_id is None and text_id is None:45        gt = torch.arange(logits.shape[0], device=logits.device)46        loss1 = F.cross_entropy(logits, gt)47        loss2 = F.cross_entropy(logits.t(), gt)        48    else:49        image_id = all_gather_grad(image_id)50        text_id = all_gather_grad(text_id)51 52        gt_image = image_id.reshape((-1, 1)) == image_id.reshape((1, -1))53        gt_text = text_id.reshape((-1, 1)) == text_id.reshape((1, -1))54        gt = torch.logical_or(gt_image, gt_text)55 56        loss1 = -torch.sum(gt * F.log_softmax(logits, dim=1)) / gt.sum()57        loss2 = -torch.sum(gt.t() * F.log_softmax(logits.t(), dim=1)) / gt.sum()58 59    return (loss1 + loss2) / 2 * get_world_size() # scale it up by the number of GPUs60