| 11 | |
| 12 | |
| 13 | class PatchNCELoss(nn.Module): |
| 14 | def __init__(self, opt): |
| 15 | super().__init__() |
| 16 | self.opt = opt |
| 17 | self.cross_entropy_loss = torch.nn.CrossEntropyLoss(reduction='none') |
| 18 | self.mask_dtype = torch.uint8 if version.parse(torch.__version__) < version.parse('1.2.0') else torch.bool |
| 19 | self.similarity_function = self._get_similarity_function() |
| 20 | self.cos = torch.nn.CosineSimilarity(dim=-1) |
| 21 | |
| 22 | def _get_similarity_function(self): |
| 23 | |
| 24 | self._cosine_similarity = torch.nn.CosineSimilarity(dim=-1) |
| 25 | return self._cosine_simililarity |
| 26 | |
| 27 | def _cosine_simililarity(self, x, y): |
| 28 | # x shape: (N, 1, C) |
| 29 | # y shape: (1, M, C) |
| 30 | # v shape: (N, M) |
| 31 | v = self._cosine_similarity(x.unsqueeze(1), y.unsqueeze(0)) |
| 32 | return v |
| 33 | |
| 34 | def forward(self, feat_q, feat_k): |
| 35 | batchSize = feat_q.shape[0] |
| 36 | feat_k = feat_k.detach() |
| 37 | l_pos = self.cos(feat_q,feat_k) |
| 38 | l_pos = l_pos.view(batchSize, 1) |
| 39 | l_neg_curbatch = self.similarity_function(feat_q.view(batchSize,1,-1),feat_k.view(1,batchSize,-1)) |
| 40 | l_neg_curbatch = l_neg_curbatch.view(1,batchSize,-1) |
| 41 | # diagonal entries are similarity between same features, and hence meaningless. |
| 42 | # just fill the diagonal with very small number, which is exp(-10) and almost zero |
| 43 | diagonal = torch.eye(batchSize, device=feat_q.device, dtype=self.mask_dtype)[None, :, :] |
| 44 | l_neg_curbatch.masked_fill_(diagonal, -10.0) |
| 45 | l_neg = l_neg_curbatch.view(-1, batchSize) |
| 46 | out = torch.cat((l_pos, l_neg), dim=1) / self.opt.nce_T |
| 47 | loss = self.cross_entropy_loss(out, torch.zeros(out.size(0), dtype=torch.long, |
| 48 | device=feat_q.device)) |
| 49 | return loss |
| 50 | |
| 51 | |
| 52 | # Used in vanilla CUT |