(self, classnames, clip_model)
| 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) |
no test coverage detected