| 214 | |
| 215 | |
| 216 | def forward(self, image): |
| 217 | |
| 218 | b = image.shape[0] |
| 219 | prompts, vision_prompts = self.prompt_learner() |
| 220 | tokenized_prompts = self.tokenized_prompts |
| 221 | image_features = self.image_encoder(image.type(self.dtype), vision_prompts) |
| 222 | image_feature_pool = image_features.mean(dim=0) |
| 223 | M = image_features.shape[0] |
| 224 | self.d = image_features.shape[-1] |
| 225 | |
| 226 | if self.dataset == 'ImageNet': |
| 227 | text_features = self.text_encoder(prompts.to(self.device1), tokenized_prompts.to(self.device1)) |
| 228 | text_features = text_features.to(self.device) |
| 229 | text_features = text_features.contiguous().view(self.N, self.n_cls, self.d) |
| 230 | text_feature_pool = text_features.mean(dim=0) |
| 231 | else: |
| 232 | text_features = self.text_encoder(prompts, tokenized_prompts).contiguous().view(self.N, self.n_cls, self.d) |
| 233 | text_feature_pool = text_features.mean(dim=0) |
| 234 | |
| 235 | image_features = F.normalize(image_features, dim=2) # N c d |
| 236 | image_feature_pool = F.normalize(image_feature_pool, dim=1) |
| 237 | text_features = F.normalize(text_features, dim=2) |
| 238 | text_feature_pool = F.normalize(text_feature_pool, dim=1) |
| 239 | sim = torch.einsum('mbd,ncd->mnbc', image_features, text_features).contiguous() |
| 240 | sim = sim.view(M,self.N,b*self.n_cls) |
| 241 | sim = sim.permute(2,0,1) |
| 242 | wdist = 1.0 - sim |
| 243 | xx=torch.zeros(b*self.n_cls, M, dtype=sim.dtype, device=sim.device).fill_(1. / M) |
| 244 | yy=torch.zeros(b*self.n_cls, self.N, dtype=sim.dtype, device=sim.device).fill_(1. / self.N) |
| 245 | |
| 246 | with torch.no_grad(): |
| 247 | KK = torch.exp(-wdist / self.eps) |
| 248 | T = self.Sinkhorn(KK,xx,yy) |
| 249 | if torch.isnan(T).any(): |
| 250 | return None |
| 251 | |
| 252 | sim_op = torch.sum(T * sim, dim=(1, 2)) |
| 253 | sim_op = sim_op.contiguous().view(b,self.n_cls) |
| 254 | |
| 255 | |
| 256 | logit_scale = self.logit_scale.exp() |
| 257 | logits = logit_scale * image_feature_pool @ text_feature_pool.t() |
| 258 | logits2 = logit_scale * sim_op |
| 259 | if self.tradeoff: |
| 260 | logits2 = logits + logits2 |
| 261 | |
| 262 | return logits2 |
| 263 | |
| 264 | |
| 265 | @TRAINER_REGISTRY.register() |