MoCo-style memory for two modalities, e.g. in CMC
| 89 | |
| 90 | |
| 91 | class CMCMoCo(BaseMoCo): |
| 92 | """MoCo-style memory for two modalities, e.g. in CMC""" |
| 93 | def __init__(self, n_dim, K=65536, T=0.07): |
| 94 | super(CMCMoCo, self).__init__(K, T) |
| 95 | # create memory queue |
| 96 | self.register_buffer('memory_1', torch.randn(K, n_dim)) |
| 97 | self.register_buffer('memory_2', torch.randn(K, n_dim)) |
| 98 | self.memory_1 = F.normalize(self.memory_1) |
| 99 | self.memory_2 = F.normalize(self.memory_2) |
| 100 | |
| 101 | def forward(self, q1, k1, q2, k2, |
| 102 | q1_jig=None, q2_jig=None, |
| 103 | all_k1=None, all_k2=None): |
| 104 | """ |
| 105 | Args: |
| 106 | q1: q of modal 1 |
| 107 | k1: k of modal 1 |
| 108 | q2: q of modal 2 |
| 109 | k2: k of modal 2 |
| 110 | q1_jig: q jig of modal 1 |
| 111 | q2_jig: q jig of modal 2 |
| 112 | all_k1: gather of k1 across nodes; otherwise use k1 |
| 113 | all_k2: gather of k2 across nodes; otherwise use k2 |
| 114 | """ |
| 115 | bsz = q1.size(0) |
| 116 | k1 = k1.detach() |
| 117 | k2 = k2.detach() |
| 118 | |
| 119 | # compute logit |
| 120 | queue1 = self.memory_1.clone().detach() |
| 121 | queue2 = self.memory_2.clone().detach() |
| 122 | logits1 = self._compute_logit(q1, k2, queue2) |
| 123 | logits2 = self._compute_logit(q2, k1, queue1) |
| 124 | if (q1_jig is not None) and (q2_jig is not None): |
| 125 | logits1_jig = self._compute_logit(q1_jig, k2, queue2) |
| 126 | logits2_jig = self._compute_logit(q2_jig, k1, queue1) |
| 127 | |
| 128 | # set label |
| 129 | labels = torch.zeros(bsz, dtype=torch.long).cuda() |
| 130 | |
| 131 | # update memory |
| 132 | all_k1 = all_k1 if all_k1 is not None else k1 |
| 133 | all_k2 = all_k2 if all_k2 is not None else k2 |
| 134 | assert all_k1.size(0) == all_k2.size(0) |
| 135 | self._update_memory(all_k1, self.memory_1) |
| 136 | self._update_memory(all_k2, self.memory_2) |
| 137 | self._update_pointer(all_k1.size(0)) |
| 138 | |
| 139 | if (q1_jig is not None) and (q2_jig is not None): |
| 140 | return logits1, logits2, logits1_jig, logits2_jig, labels |
| 141 | else: |
| 142 | return logits1, logits2, labels |
nothing calls this directly
no outgoing calls
no test coverage detected