Test for https://gist.github.com/KeremTurgutlu/4a6f7078dc62f292c85b9903197c75f7

basic_batch_all_gather.py

import os
import torch
import torch.distributed as dist
from torch.multiprocessing import Process
from torchvision import datasets, transforms
import numpy as np

import  torch.nn as nn
import torch.nn.functional as F
import random

def set_seed(s, reproducible=False):
    "Set random seed for `random`, `torch`, and `numpy` (where available)"
    try: torch.manual_seed(s)
    except NameError: pass
    try: torch.cuda.manual_seed_all(s)
    except NameError: pass
    try: np.random.seed(s%(2**32-1))
    except NameError: pass
    random.seed(s)
    if reproducible:
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False

set_seed(42)
image_inputs = torch.randn(8, 4)
set_seed(42)
text_inputs  = torch.randn(8, 4)
set_seed(42)
image_encoder = nn.Linear(4, 2, bias=False)
set_seed(42)
text_encoder  = nn.Linear(4, 2, bias=False)

# print(image_inputs)
# print(text_inputs)

# print(image_encoder.weight)
# print(text_encoder.weight)


image_embeddings = image_encoder(image_inputs)
text_embeddings  = text_encoder(text_inputs)

image_embeddings = F.normalize(image_embeddings)
text_embeddings = F.normalize(text_embeddings)

image_embeddings.shape, text_embeddings.shape

cosine_sim = image_embeddings @ text_embeddings.T

loss = F.cross_entropy(cosine_sim, torch.arange(len(cosine_sim)), reduction="none")
print(loss)

loss = loss.mean()
print(loss)

loss.backward()
print(image_encoder.weight.grad, text_encoder.weight.grad)


# print(image_encoder.weight)
# print(text_encoder.weight)

Following code results in:

Loss: 1.393433690071106
image_encoder.weight.grad: tensor([[-0.0559,  0.0055, -0.0663,  0.0330],
        [ 0.1006, -0.0770,  0.0099,  0.0591]])
text_encoder.weight.grad: tensor([[-0.0575,  0.0046, -0.0505,  0.0319],
        [ 0.0871, -0.0862,  0.0119,  0.0704]])