Multimodal: DALL·E 2, Imagen, Stable Diffusion

import torch, torch.nn as nn

torch.manual_seed(0)
# Toy paired data: 64 "concepts", each seen as an image vector and a text vector
concepts = torch.randn(64, 16)
images = concepts @ torch.randn(16, 32) + 0.1 * torch.randn(64, 32)
texts = concepts @ torch.randn(16, 24) + 0.1 * torch.randn(64, 24)

img_enc, txt_enc = nn.Linear(32, 16), nn.Linear(24, 16)
opt = torch.optim.Adam([*img_enc.parameters(), *txt_enc.parameters()], lr=1e-2)

def retrieval_acc():
    with torch.no_grad():
        i = nn.functional.normalize(img_enc(images), dim=1)
        t = nn.functional.normalize(txt_enc(texts), dim=1)
        return ((i @ t.T).argmax(1) == torch.arange(64)).float().mean().item()

print(f"before training: text->image match {retrieval_acc():.0%}")
for step in range(300):
    i = nn.functional.normalize(img_enc(images), dim=1)
    t = nn.functional.normalize(txt_enc(texts), dim=1)
    logits = i @ t.T / 0.07                                # similarity of every pair
    labels = torch.arange(64)                              # matching pairs are on the diagonal
    loss = (nn.functional.cross_entropy(logits, labels) +
            nn.functional.cross_entropy(logits.T, labels)) / 2   # CLIP contrastive loss
    opt.zero_grad(); loss.backward(); opt.step()
print(f"after training:  text->image match {retrieval_acc():.0%}  (loss {loss.item():.3f})")