File size: 1,585 Bytes
803ef9e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 |
from functools import partial
import torch
import torch.nn as nn
import torch.nn.functional as F
from .base import BaseMethod
def contrastive_loss(x0, x1, tau, norm):
# https://github.com/google-research/simclr/blob/master/objective.py
bsize = x0.shape[0]
target = torch.arange(bsize).cuda()
eye_mask = torch.eye(bsize).cuda() * 1e9
if norm:
x0 = F.normalize(x0, p=2, dim=1)
x1 = F.normalize(x1, p=2, dim=1)
logits00 = x0 @ x0.t() / tau - eye_mask
logits11 = x1 @ x1.t() / tau - eye_mask
logits01 = x0 @ x1.t() / tau
logits10 = x1 @ x0.t() / tau
return (
F.cross_entropy(torch.cat([logits01, logits00], dim=1), target)
+ F.cross_entropy(torch.cat([logits10, logits11], dim=1), target)
) / 2
class Contrastive(BaseMethod):
""" implements contrastive loss https://arxiv.org/abs/2002.05709 """
def __init__(self, cfg):
""" init additional BN used after head """
super().__init__(cfg)
self.bn_last = nn.BatchNorm1d(cfg.emb)
self.loss_f = partial(contrastive_loss, tau=cfg.tau, norm=cfg.norm)
def forward(self, samples):
bs = len(samples[0])
h = [self.model(x.cuda(non_blocking=True)) for x in samples]
h = self.bn_last(self.head(torch.cat(h)))
loss = 0
for i in range(len(samples) - 1):
for j in range(i + 1, len(samples)):
x0 = h[i * bs : (i + 1) * bs]
x1 = h[j * bs : (j + 1) * bs]
loss += self.loss_f(x0, x1)
loss /= self.num_pairs
return loss
|