diff options
| -rw-r--r-- | libs/criteria.py | 25 | ||||
| -rw-r--r-- | readme.md | 2 |
2 files changed, 26 insertions, 1 deletions
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 @@ -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 |