| 228 | return T |
| 229 | |
| 230 | def forward(self, image): |
| 231 | |
| 232 | b = image.shape[0] |
| 233 | image_features = self.image_encoder(image.type(self.dtype)) |
| 234 | image_feature_pool = image_features[0] |
| 235 | image_features = image_features[1:] |
| 236 | M = image_features.shape[0] |
| 237 | self.d = image_features.shape[-1] |
| 238 | |
| 239 | prompts = self.prompt_learner() |
| 240 | tokenized_prompts = self.tokenized_prompts |
| 241 | if self.dataset == "ImageNet": |
| 242 | text_features = self.text_encoder(prompts.to(self.device1), tokenized_prompts.to(self.device1)) |
| 243 | text_features = text_features.to(self.device) |
| 244 | text_features = text_features.contiguous().view(self.N, self.n_cls, self.d) |
| 245 | text_feature_pool = text_features.mean(dim=0) |
| 246 | else: |
| 247 | text_features = self.text_encoder(prompts, tokenized_prompts) |
| 248 | text_features = text_features.contiguous().view(self.N, self.n_cls, self.d) |
| 249 | text_feature_pool = text_features.mean(dim=0) |
| 250 | |
| 251 | |
| 252 | image_features = F.normalize(image_features, dim=2) |
| 253 | image_feature_pool = F.normalize(image_feature_pool, dim=1) |
| 254 | text_features = F.normalize(text_features, dim=2) |
| 255 | text_feature_pool = F.normalize(text_feature_pool, dim=1) |
| 256 | |
| 257 | sim = torch.einsum('mbd,ncd->mnbc', image_features, text_features).contiguous() |
| 258 | sim = sim.view(M,self.N,b*self.n_cls) |
| 259 | sim = sim.permute(2,0,1) |
| 260 | wdist = 1.0 - sim |
| 261 | xx=torch.zeros(b*self.n_cls, M, dtype=sim.dtype, device=sim.device).fill_(1. / M) |
| 262 | yy=torch.zeros(b*self.n_cls, self.N, dtype=sim.dtype, device=sim.device).fill_(1. / self.N) |
| 263 | |
| 264 | |
| 265 | |
| 266 | with torch.no_grad(): |
| 267 | KK = torch.exp(-wdist / self.eps) |
| 268 | T = self.Sinkhorn(KK,xx,yy) |
| 269 | if torch.isnan(T).any(): |
| 270 | return None |
| 271 | |
| 272 | |
| 273 | sim_op = torch.sum(T * sim, dim=(1, 2)) |
| 274 | sim_op = sim_op.contiguous().view(b,self.n_cls) |
| 275 | |
| 276 | |
| 277 | logit_scale = self.logit_scale.exp() |
| 278 | logits = logit_scale * image_feature_pool @ text_feature_pool.t() |
| 279 | logits2 = logit_scale * sim_op |
| 280 | if self.dataset == "ImageNet": |
| 281 | logits2 = (logits2 + logits) |
| 282 | return logits2 |
| 283 | |
| 284 | |
| 285 | @TRAINER_REGISTRY.register() |