import torch
import torch.nn as nn
import torch.nn.functional as F
class CLIPContrastiveLoss(nn.Module):
def __init__(self, init_tau: float = 0.07):
super().__init__()
# Learnable log temperature parameter log(1/tau)
self.logit_scale = nn.Parameter(torch.tensor(torch.log(torch.tensor(1.0 / init_tau))))
def forward(self, vision_embeds: torch.Tensor, text_embeds: torch.Tensor):
# L2 Normalization
v_norm = F.normalize(vision_embeds, p=2, dim=-1)
t_norm = F.normalize(text_embeds, p=2, dim=-1)
# Cosine similarity matrix scaled by exp(logit_scale)
logit_scale = torch.exp(self.logit_scale)
logits_per_image = logit_scale * (v_norm @ t_norm.T) # [N, N]
logits_per_text = logits_per_image.T
# Ground truth diagonal targets
labels = torch.arange(vision_embeds.shape[0], device=vision_embeds.device)
# Symmetric Cross-Entropy Loss
loss_v = F.cross_entropy(logits_per_image, labels)
loss_t = F.cross_entropy(logits_per_text, labels)
return (loss_v + loss_t) / 2.0