From 2cc459e3e4b2d559b5d8aa757c694db02ccd0e2a Mon Sep 17 00:00:00 2001 From: Jordan Gong Date: Thu, 14 Jul 2022 17:41:49 +0800 Subject: Implement InfoNCE loss --- libs/criteria.py | 25 +++++++++++++++++++++++++ 1 file changed, 25 insertions(+) create mode 100644 libs/criteria.py (limited to 'libs') diff --git a/libs/criteria.py b/libs/criteria.py new file mode 100644 index 0000000..6954cf3 --- /dev/null +++ b/libs/criteria.py @@ -0,0 +1,25 @@ +import torch +from torch import nn, Tensor +from torch.nn import functional as F + + +class InfoNCELoss(nn.Module): + def __init__(self, temp=0.01): + super().__init__() + self.temp = temp + + def forward(self, feat1: Tensor, feat2: Tensor) -> tuple[Tensor, Tensor]: + bz = feat1.size(0) + feat1_norm = F.normalize(feat1) + feat2_norm = F.normalize(feat2) + logits = feat1_norm @ feat2_norm.T + pos_logits_mask = torch.eye(bz, dtype=torch.bool) + pos_logits = logits[pos_logits_mask].unsqueeze(-1) + neg_logits = logits[~pos_logits_mask].view(bz, -1) + # Put the positive at first (0-th) and maximize its likelihood + logits = torch.cat([pos_logits, neg_logits], dim=1) + labels = torch.zeros(bz, dtype=torch.long, device=feat1.device) + loss_contra = F.cross_entropy(logits / self.temp, labels) + acc_contra = (logits.argmax(dim=1) == labels).float().mean() + + return loss_contra, acc_contra -- cgit v1.2.3