| 210 | |
| 211 | |
| 212 | class CustomCLIP(nn.Module): |
| 213 | def __init__(self, classnames, clip_model): |
| 214 | super().__init__() |
| 215 | self.n_cls = len(classnames) |
| 216 | self.prompt_learner = PromptLearner(classnames, clip_model) |
| 217 | self.tokenized_prompts = self.prompt_learner.tokenized_prompts |
| 218 | self.image_encoder = clip_model.visual |
| 219 | self.device0 = torch.device("cuda:0") |
| 220 | self.device = torch.device("cuda") |
| 221 | self.text_encoder = TextEncoder(clip_model) |
| 222 | self.logit_scale = clip_model.logit_scale |
| 223 | self.dtype = clip_model.dtype |
| 224 | self.N = 4 #cfg.MODEL.N |
| 225 | self.use_uniform = True |
| 226 | self.eps = 0.1 |
| 227 | self.max_iter = 100 |
| 228 | |
| 229 | def Sinkhorn(self, K, u, v): |
| 230 | r = torch.ones_like(u) |
| 231 | c = torch.ones_like(v) |
| 232 | thresh = 1e-2 |
| 233 | for i in range(self.max_iter): |
| 234 | r0 = r |
| 235 | r = u / torch.matmul(K, c.unsqueeze(-1)).squeeze(-1) |
| 236 | c = v / torch.matmul(K.permute(0, 2, 1).contiguous(), r.unsqueeze(-1)).squeeze(-1) |
| 237 | err = (r - r0).abs().mean() |
| 238 | if err.item() < thresh: |
| 239 | break |
| 240 | |
| 241 | T = torch.matmul(r.unsqueeze(-1), c.unsqueeze(-2)) * K |
| 242 | |
| 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 | |