| 243 | return T |
| 244 | |
| 245 | def forward(self, image): |
| 246 | |
| 247 | b = image.shape[0] |
| 248 | image_features = self.image_encoder(image.type(self.dtype)) |
| 249 | image_feature_pool = image_features[0] |
| 250 | image_features = image_features[1:] |
| 251 | M = image_features.shape[0] |
| 252 | self.d = image_features.shape[-1] |
| 253 | |
| 254 | prompts = self.prompt_learner() |
| 255 | |
| 256 | tokenized_prompts = self.tokenized_prompts |
| 257 | |
| 258 | |
| 259 | text_features = self.text_encoder(prompts.to(self.device), tokenized_prompts.to(self.device)) |
| 260 | text_features = text_features.to(self.device0) |
| 261 | text_features = text_features.contiguous().view(self.N, self.n_cls, self.d) |
| 262 | text_feature_pool = text_features.mean(dim=0) |
| 263 | |
| 264 | |
| 265 | image_features = F.normalize(image_features, dim=2) |
| 266 | image_feature_pool = F.normalize(image_feature_pool, dim=1) |
| 267 | text_features = F.normalize(text_features, dim=2) |
| 268 | text_feature_pool = F.normalize(text_feature_pool, dim=1) |
| 269 | |
| 270 | sim = torch.einsum('mbd,ncd->mnbc', image_features, text_features).contiguous() |
| 271 | sim = sim.view(M,self.N,b*self.n_cls) |
| 272 | sim = sim.permute(2,0,1) |
| 273 | wdist = 1.0 - sim |
| 274 | |
| 275 | xx=torch.zeros(b*self.n_cls, M, dtype=sim.dtype, device=sim.device).fill_(1. / M) |
| 276 | yy=torch.zeros(b*self.n_cls, self.N, dtype=sim.dtype, device=sim.device).fill_(1. / self.N) |
| 277 | |
| 278 | with torch.no_grad(): |
| 279 | KK = torch.exp(-wdist / self.eps) |
| 280 | T = self.Sinkhorn(KK,xx,yy) |
| 281 | try: |
| 282 | torch.isnan(T).any() |
| 283 | except None: |
| 284 | print('There is none value in your tensor, please try to adjust #thre and #eps to align data.') |
| 285 | |
| 286 | |
| 287 | sim_op = torch.sum(T * sim, dim=(1, 2)) |
| 288 | sim_op = sim_op.contiguous().view(b,self.n_cls) |
| 289 | |
| 290 | logit_scale = self.logit_scale.exp() |
| 291 | logits = logit_scale * image_feature_pool @ text_feature_pool.t() |
| 292 | |
| 293 | logits2 = logit_scale * sim_op |
| 294 | logits2 = (0.5*logits2 + 0.5*logits) |
| 295 | return logits2, image_feature_pool,text_feature_pool, image_features |
| 296 | |
| 297 | |
| 298 | def get_arguments(): |