| 51 | |
| 52 | # Used in vanilla CUT |
| 53 | class PatchNCELoss2(nn.Module): |
| 54 | def __init__(self, opt): |
| 55 | super().__init__() |
| 56 | self.opt = opt |
| 57 | self.cross_entropy_loss = torch.nn.CrossEntropyLoss(reduction='none') |
| 58 | self.mask_dtype = torch.uint8 if version.parse(torch.__version__) < version.parse('1.2.0') else torch.bool |
| 59 | |
| 60 | def forward(self, feat_q, feat_k): |
| 61 | batchSize = feat_q.shape[0] |
| 62 | dim = feat_q.shape[1] |
| 63 | feat_k = feat_k.detach() |
| 64 | |
| 65 | # pos logit |
| 66 | l_pos = torch.bmm(feat_q.view(batchSize, 1, -1), feat_k.view(batchSize, -1, 1)) |
| 67 | l_pos = l_pos.view(batchSize, 1) |
| 68 | |
| 69 | # neg logit |
| 70 | |
| 71 | # Should the negatives from the other samples of a minibatch be utilized? |
| 72 | # In CUT and FastCUT, we found that it's best to only include negatives |
| 73 | # from the same image. Therefore, we set |
| 74 | # --nce_includes_all_negatives_from_minibatch as False |
| 75 | # However, for single-image translation, the minibatch consists of |
| 76 | # crops from the "same" high-resolution image. |
| 77 | # Therefore, we will include the negatives from the entire minibatch. |
| 78 | if self.opt.nce_includes_all_negatives_from_minibatch: |
| 79 | # reshape features as if they are all negatives of minibatch of size 1. |
| 80 | batch_dim_for_bmm = 1 |
| 81 | else: |
| 82 | batch_dim_for_bmm = self.opt.batch_size |
| 83 | |
| 84 | # reshape features to batch size |
| 85 | feat_q = feat_q.view(batch_dim_for_bmm, -1, dim) |
| 86 | feat_k = feat_k.view(batch_dim_for_bmm, -1, dim) |
| 87 | npatches = feat_q.size(1) |
| 88 | l_neg_curbatch = torch.bmm(feat_q, feat_k.transpose(2, 1)) |
| 89 | |
| 90 | # diagonal entries are similarity between same features, and hence meaningless. |
| 91 | # just fill the diagonal with very small number, which is exp(-10) and almost zero |
| 92 | diagonal = torch.eye(npatches, device=feat_q.device, dtype=self.mask_dtype)[None, :, :] |
| 93 | l_neg_curbatch.masked_fill_(diagonal, -10.0) |
| 94 | l_neg = l_neg_curbatch.view(-1, npatches) |
| 95 | |
| 96 | out = torch.cat((l_pos, l_neg), dim=1) / self.opt.nce_T |
| 97 | |
| 98 | loss = self.cross_entropy_loss(out, torch.zeros(out.size(0), dtype=torch.long, |
| 99 | device=feat_q.device)) |
| 100 | return loss |
| 101 | |