| 194 | |
| 195 | |
| 196 | class CustomCLIP(nn.Module): |
| 197 | def __init__(self, cfg, classnames, clip_model): |
| 198 | super().__init__() |
| 199 | self.n_cls = len(classnames) |
| 200 | self.prompt_learner = PromptLearner(cfg, classnames, clip_model) |
| 201 | self.tokenized_prompts = self.prompt_learner.tokenized_prompts |
| 202 | self.image_encoder = clip_model.visual |
| 203 | self.text_encoder = TextEncoder(clip_model) |
| 204 | self.logit_scale = clip_model.logit_scale |
| 205 | self.dtype = clip_model.dtype |
| 206 | self.device = torch.device("cuda:0") |
| 207 | self.device1 = torch.device("cuda") |
| 208 | self.N = cfg.TRAINER.PLOT.N |
| 209 | self.dataset = cfg.DATASET.NAME |
| 210 | self.use_uniform = True |
| 211 | self.eps = 0.1 |
| 212 | self.max_iter = 100 |
| 213 | |
| 214 | def Sinkhorn(self, K, u, v): |
| 215 | r = torch.ones_like(u) |
| 216 | c = torch.ones_like(v) |
| 217 | thresh = 1e-2 |
| 218 | for i in range(self.max_iter): |
| 219 | r0 = r |
| 220 | r = u / torch.matmul(K, c.unsqueeze(-1)).squeeze(-1) |
| 221 | c = v / torch.matmul(K.permute(0, 2, 1).contiguous(), r.unsqueeze(-1)).squeeze(-1) |
| 222 | err = (r - r0).abs().mean() |
| 223 | if err.item() < thresh: |
| 224 | break |
| 225 | |
| 226 | T = torch.matmul(r.unsqueeze(-1), c.unsqueeze(-2)) * K |
| 227 | |
| 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) |