Input: im_q: a batch of query images im_k: a batch of key images queue: a queue from which to pick negative samples Output: logits, targets
(self, img_q, img_k, queue)
| 152 | return x_gather[idx_this] |
| 153 | |
| 154 | def forward(self, img_q, img_k, queue): |
| 155 | """ |
| 156 | Input: |
| 157 | im_q: a batch of query images |
| 158 | im_k: a batch of key images |
| 159 | queue: a queue from which to pick negative samples |
| 160 | Output: |
| 161 | logits, targets |
| 162 | """ |
| 163 | |
| 164 | # compute query features |
| 165 | q = self.encoder_q(img_q) # queries: NxC |
| 166 | q = nn.functional.normalize(q, dim=1) |
| 167 | |
| 168 | # compute key features |
| 169 | with torch.no_grad(): # no gradient to keys |
| 170 | # shuffle for making use of BN |
| 171 | if self._use_ddp_or_ddp2(self.trainer): |
| 172 | img_k, idx_unshuffle = self._batch_shuffle_ddp(img_k) |
| 173 | |
| 174 | k = self.encoder_k(img_k) # keys: NxC |
| 175 | k = nn.functional.normalize(k, dim=1) |
| 176 | |
| 177 | # undo shuffle |
| 178 | if self._use_ddp_or_ddp2(self.trainer): |
| 179 | k = self._batch_unshuffle_ddp(k, idx_unshuffle) |
| 180 | |
| 181 | # compute logits |
| 182 | # Einstein sum is more intuitive |
| 183 | # positive logits: Nx1 |
| 184 | l_pos = torch.einsum("nc,nc->n", [q, k]).unsqueeze(-1) |
| 185 | # negative logits: NxK |
| 186 | l_neg = torch.einsum("nc,ck->nk", [q, queue.clone().detach()]) |
| 187 | |
| 188 | # logits: Nx(1+K) |
| 189 | logits = torch.cat([l_pos, l_neg], dim=1) |
| 190 | |
| 191 | # apply temperature |
| 192 | logits /= self.hparams.softmax_temperature |
| 193 | |
| 194 | # labels: positive key indicators |
| 195 | labels = torch.zeros(logits.shape[0], dtype=torch.long) |
| 196 | labels = labels.type_as(logits) |
| 197 | |
| 198 | return logits, labels, k |
| 199 | |
| 200 | def training_step(self, batch, batch_idx): |
| 201 | # in STL10 we pass in both lab+unl for online ft |
nothing calls this directly
no test coverage detected