EMBED_DIM = 512
TRANSFORMER_EMBED_DIM = 768
MAX_LEN = 128 # Maximum length of text
TEXT_MODEL = "distilbert-base-multilingual-cased" EPOCHS = 5
BATCH_SIZE = 64

img = inv_tfm(img)
plt.imshow(np.rot90(img.transpose(0, 2), 3))
plt.title(target)
plt.show()

train_len = int(0.8*len(cap))
train_data, valid_data = random_split(cap, [train_len, len(cap) - train_len])
train_dl = DataLoader(train_data, BATCH_SIZE, pin_memory=True, shuffle=True, num_workers=4, drop_last=True)
valid_dl = DataLoader(valid_data, BATCH_SIZE, pin_memory=True, shuffle=False, num_workers=4, drop_last=False)

class Projection(nn.Module): def __init__(self, d_in: int, d_out: int, p: float=0.5) -> None: super().__init__() self.linear1 = nn.Linear(d_in, d_out, bias=False) self.linear2 = nn.Linear(d_out, d_out, bias=False) self.layer_norm = nn.LayerNorm(d_out) self.drop = nn.Dropout(p) def forward(self, x: torch.Tensor) -> torch.Tensor: embed1 = self.linear1(x) embed2 = self.drop(self.linear2(F.gelu(embed1))) embeds = self.layer_norm(embed1 + embed2) return embeds

class VisionEncoder(nn.Module): def __init__(self, d_out: int) -> None: super().__init__() base = models.resnet34(pretrained=True) d_in = base.fc.in_features base.fc = nn.Identity() self.base = base self.projection = Projection(d_in, d_out) for p in self.base.parameters(): p.requires_grad = False def forward(self, x): projected_vec = self.projection(self.base(x)) projection_len = torch.norm(projected_vec, dim=-1, keepdim=True) return projected_vec / projection_len

class TextEncoder(nn.Module): def __init__(self, d_out: int) -> None: super().__init__() self.base = AutoModel.from_pretrained(TEXT_MODEL) self.projection = Projection(TRANSFORMER_EMBED_DIM, d_out) for p in self.base.parameters(): p.requires_grad = False def forward(self, x): out = self.base(**x)[0] out = out[:, 0, :] # get CLS token output projected_vec = self.projection(out) projection_len = torch.norm(projected_vec, dim=-1, keepdim=True) return projected_vec / projection_len class Tokenizer: def __init__(self, tokenizer: BertTokenizer) -> None: self.tokenizer = tokenizer def __call__(self, x: str) -> AutoTokenizer: return self.tokenizer( x, max_length=MAX_LEN, truncation=True, padding=True, return_tensors="pt" )

def contrastive_loss(logits, dim): neg_ce = torch.diag(F.log_softmax(logits, dim=dim)) return -neg_ce.mean() def clip_loss(similarity: torch.Tensor) -> torch.Tensor: caption_loss = contrastive_loss(similarity, dim=0) image_loss = contrastive_loss(similarity, dim=1) return (caption_loss + image_loss) / 2.0 def metrics(similarity: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: y = torch.arange(len(similarity)).to(similarity.device) img2cap_match_idx = similarity.argmax(dim=1) cap2img_match_idx = similarity.argmax(dim=0) img_acc = (img2cap_match_idx == y).float().mean() cap_acc = (cap2img_match_idx == y).float().mean() return img_acc, cap_acc

class Model(pl.LightningModule): def __init__(self, lr: float = 1e-3 ) -> None: super().__init__() self.vision_encoder = VisionEncoder(EMBED_DIM) self.caption_encoder = TextEncoder(EMBED_DIM) self.tokenizer = Tokenizer(AutoTokenizer.from_pretrained(TEXT_MODEL)) self.lr = lr def common_step(self, batch: Tuple[torch.Tensor, List[str]]) -> torch.Tensor: images, text = batch device = images.device text_dev = {k: v.to(device) for k, v in self.tokenizer(text).items()} image_embed = self.vision_encoder(images) caption_embed = self.caption_encoder(text_dev) similarity = caption_embed @ image_embed.T loss = clip_loss(similarity) img_acc, cap_acc = metrics(similarity) return loss, img_acc, cap_acc def training_step( self, batch: Tuple[torch.Tensor, List[str]], *args: list ) -> torch.Tensor: loss, img_acc, cap_acc = self.common_step(batch) self.log("training_loss", loss, on_step=True) self.log("training_img_acc", img_acc, on_step=True, prog_bar=True) self.log("training_cap_acc", cap_acc, on_step=True, prog_bar=True) return loss def validation_step( self, batch: Tuple[torch.Tensor, List[str]], *args: list ) -> torch.Tensor: loss, img_acc, cap_acc = self.common_step(batch) self.log("validation_loss", loss, on_step=True) self.log("validation_img_acc", img_acc, on_step=True, prog_bar=True) self.log("validation_cap_acc", cap_acc, on_step=True, prog_bar=True) return loss def configure_optimizers(self) -> torch.optim.Optimizer: vision_params = {"params": self.vision_encoder.projection.parameters(), "lr": self.lr} caption_params = {"params": self.caption_encoder.projection.parameters() , "lr": self.lr} return torch.optim.Adam([vision_params, caption_params])

model = Model(1e-3)
trainer = pl.Trainer( max_epochs= EPOCHS, gpus=torch.cuda.device_count(), gradient_clip_val=1.0, precision=16
)
trainer.fit(model, train_dl, valid_dl)

Downloading: "https://download.pytorch.org/models/resnet34-333f7ec4.pth" to /root/.cache/torch/hub/checkpoints/resnet34-333f7ec4.pth GPU available: True, used: True TPU available: None, using: 0 TPU cores Using native 16bit precision. | Name | Type | Params
--------------------------------------------------
0 | vision_encoder | VisionEncoder | 21.8 M
1 | caption_encoder | TextEncoder | 135 M --------------------------------------------------
1.2 M Trainable params
156 M Non-trainable params
157 M Total params
628.802 Total estimated model params size (MB) 1

%reload_ext tensorboard
%tensorboard --logdir ./lightning_logs/

similarity = caption_embed @ image_embed.T
val, closest = similarity.topk(5, dim=-1)
similarity.shape

torch.Size([64, 16557])

draw_result(2, similarity)

A baseball player in the outfield with his hands up, standing next to a team mascot.

draw_result(1, similarity)

A watch and clock repair shop window with clocks on display.

draw_result(10, similarity)

A person on a skateboard on the ground.

text = "Una cebra de pie con la cabeza gacha y comiendo hierba en el suelo de tierra."
text_dev = {k: v.to(device) for k, v in tokenizer(text).items()}
with torch.no_grad(): caption_embed_text = caption_encoder(text_dev) similarity_text = caption_embed_text @ image_embed.T

draw_result_single_query(10, similarity_text)

Skateboarder conducting a trick with bicycles in the background.

text = "Un ordinateur portable est affiché sur une petite plate-forme en bois."
text_dev = {k: v.to(device) for k, v in tokenizer(text).items()}
with torch.no_grad(): caption_embed_text = caption_encoder(text_dev) similarity_text = caption_embed_text @ image_embed.T draw_result_single_query(3, similarity_text)

Laptop computer on a small table on the side of a bed

text = "Магазин с разными часами"
text_dev = {k: v.to(device) for k, v in tokenizer(text).items()}
with torch.no_grad(): caption_embed_text = caption_encoder(text_dev) similarity_text = caption_embed_text @ image_embed.T draw_result_single_query(1, similarity_text)

A room filled with clocks through a window.

text = "bear"
text_dev = {k: v.to(device) for k, v in tokenizer(text).items()}
with torch.no_grad(): caption_embed_text = caption_encoder(text_dev) similarity_text = caption_embed_text @ image_embed.T draw_result_single_query(1, similarity_text)

Large collection of digital and analog clocks on display.