(self, input, all_gather=False)
| 175 | return output |
| 176 | |
| 177 | def forward(self, input, all_gather=False): |
| 178 | # input |
| 179 | images = input['images'] |
| 180 | texts = input['captions'] |
| 181 | texts = self.sample_captions(texts) |
| 182 | # text&image encode |
| 183 | image_features = self.encode_image(images) |
| 184 | text_features = self.text_encoder(texts) |
| 185 | |
| 186 | # print('feature mean', image_features.mean(-1), 'abs mean', image_features.abs().mean(-1), 'abs max', image_features.abs().max(), 'std', image_features.std(-1), flush=True) |
| 187 | |
| 188 | # normalized features |
| 189 | image_features = image_features / (image_features.norm(dim=-1, keepdim=True)) |
| 190 | text_features = text_features / (text_features.norm(dim=-1, keepdim=True)+1e-10) |
| 191 | |
| 192 | # cosine similarity as logits |
| 193 | logit_scale = self.logit_scale.exp() |
| 194 | logit_scale.data = torch.clamp(logit_scale.data, max=100) |
| 195 | |
| 196 | if self.training and self.use_allgather or all_gather: |
| 197 | gathered_image_features = self.all_gather(image_features) |
| 198 | gathered_text_features = self.all_gather(text_features) |
| 199 | |
| 200 | logits_per_image = logit_scale * image_features @ gathered_text_features.t() |
| 201 | logits_per_text = logit_scale * text_features @ gathered_image_features.t() |
| 202 | else: |
| 203 | logits_per_image = logit_scale * image_features @ text_features.t() |
| 204 | logits_per_text = logit_scale * text_features @ image_features.t() |
| 205 | |
| 206 | return logits_per_image, logits_per_text |
| 207 | |
| 208 | |
| 209 | class SLIP(CLIP): |
nothing calls this directly
no test coverage detected