| 175 | |
| 176 | |
| 177 | class CustomCLIP(nn.Module): |
| 178 | def __init__(self, cfg, classnames, clip_model): |
| 179 | super().__init__() |
| 180 | self.prompt_learner = MultiModalPromptLearner(cfg, classnames, clip_model) |
| 181 | self.tokenized_prompts = self.prompt_learner.tokenized_prompts |
| 182 | self.image_encoder = clip_model.visual |
| 183 | self.text_encoder = TextEncoder(clip_model) |
| 184 | self.logit_scale = clip_model.logit_scale |
| 185 | self.dtype = clip_model.dtype |
| 186 | self.N = cfg.TRAINER.PLOTPP.N |
| 187 | self.n_cls = len(classnames) |
| 188 | self.tradeoff = cfg.TRAINER.PLOTPP.TRADE_OFF # whether use OT |
| 189 | self.eps = 0.1 |
| 190 | self.max_iter = 100 |
| 191 | self.dataset = cfg.DATASET.NAME |
| 192 | if self.dataset== 'ImageNet': |
| 193 | self.device = torch.device('cuda:0') |
| 194 | self.device1 = torch.device("cuda") |
| 195 | else: |
| 196 | self.device = torch.device(cfg['DEVICE']) |
| 197 | self.device1 = torch.device("cuda") |
| 198 | |
| 199 | def Sinkhorn(self, K, u, v): |
| 200 | r = torch.ones_like(u) |
| 201 | c = torch.ones_like(v) |
| 202 | thresh = 1e-2 |
| 203 | for i in range(self.max_iter): |
| 204 | r0 = r |
| 205 | r = u / torch.matmul(K, c.unsqueeze(-1)).squeeze(-1) |
| 206 | c = v / torch.matmul(K.permute(0, 2, 1).contiguous(), r.unsqueeze(-1)).squeeze(-1) |
| 207 | err = (r - r0).abs().mean() |
| 208 | if err.item() < thresh: |
| 209 | break |
| 210 | |
| 211 | T = torch.matmul(r.unsqueeze(-1), c.unsqueeze(-2)) * K |
| 212 | |
| 213 | return T |
| 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 | |