(self, feat_q, feat_k)
| 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 |
nothing calls this directly
no outgoing calls
no test coverage detected