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 +++++++++++++++++++++++++ readme.md | 2 +- 2 files changed, 26 insertions(+), 1 deletion(-) create mode 100644 libs/criteria.py 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 diff --git a/readme.md b/readme.md index 29dc613..2d5bb74 100644 --- a/readme.md +++ b/readme.md @@ -4,7 +4,7 @@ - [x] Random color distortion (from SimCLR) - [x] Random Gaussian blur (from SimCLR) - [x] Random multi-crop (from SwAV) - - [ ] InfoNCE loss + - [x] InfoNCE loss - [ ] Momentum encoder - [ ] Selective Kernel (SK) convolution layer (used in SimCLR v2) - [x] LARS optimizer -- cgit v1.2.3