| 45 | |
| 46 | |
| 47 | class TextPrompter(nn.Module): |
| 48 | def __init__(self, classnames, args, fill=None): |
| 49 | super().__init__() |
| 50 | n_cls = len(classnames) |
| 51 | self.args = args |
| 52 | #clip_model, transform = clip.load(args.model) |
| 53 | clip_model, transform = clip.load(args.model) |
| 54 | clip_model.cpu() |
| 55 | dtype = clip_model.dtype |
| 56 | ctx_dim = clip_model.ln_final.weight.shape[0] |
| 57 | clip_imsize = clip_model.visual.input_resolution |
| 58 | cfg_imsize = 224 |
| 59 | n = 3 |
| 60 | assert cfg_imsize == clip_imsize, f"cfg_imsize ({cfg_imsize}) must equal to clip_imsize ({clip_imsize})" |
| 61 | |
| 62 | if fill is None: |
| 63 | fill = args.target if args.dataset != 'visda' else 'real' |
| 64 | fill = fill.replace('_', ' ').lower() |
| 65 | pre = 'a' |
| 66 | naive_prompt_prefix_tar = "{} {} photo of a".format(pre, fill) |
| 67 | |
| 68 | ctx_vectors = torch.empty(n_cls, n, ctx_dim, dtype=dtype) |
| 69 | |
| 70 | nn.init.normal_(ctx_vectors, std=0.02) |
| 71 | prompt_prefix = "a " + " ".join(["C"] * n) |
| 72 | |
| 73 | self.ctx = nn.Parameter(ctx_vectors) # to be optimized |
| 74 | |
| 75 | classnames = [name.replace("_", " ") for name in classnames] |
| 76 | name_lens = [len(name.split(' ')) for name in classnames] |
| 77 | naive_prompts = [naive_prompt_prefix_tar + " " + name + "." for name in classnames] |
| 78 | |
| 79 | prompts = [prompt_prefix + " photo of a " + name + "." for name in classnames] |
| 80 | print("Naive prompt: {}".format(naive_prompts[0])) |
| 81 | |
| 82 | tokenized_prompts = torch.cat([clip.tokenize(p) for p in prompts]) |
| 83 | naive_tokenized_prompts = torch.cat([clip.tokenize(p) for p in naive_prompts]) |
| 84 | |
| 85 | with torch.no_grad(): |
| 86 | embedding = clip_model.token_embedding(tokenized_prompts).type(dtype) |
| 87 | naive_embedding = clip_model.token_embedding(naive_tokenized_prompts).type(dtype) |
| 88 | |
| 89 | # These token vectors will be saved when in save_model(), |
| 90 | # but they should be ignored in load_model() as we want to use |
| 91 | # those computed using the current class names |
| 92 | #tokenized_prompts = torch.cat([tokenized_prompts, naive_tokenized_prompts]) |
| 93 | self.register_buffer("token_prefix", embedding[:, :1, :]) # SOS |
| 94 | self.register_buffer("token_suffix", embedding[:, 1 + n:, :]) # CLS, EOS |
| 95 | |
| 96 | self.n_cls = n_cls |
| 97 | self.csc = True |
| 98 | self.tokenized_prompts = tokenized_prompts |
| 99 | self.naive_tokenized_prompts = naive_tokenized_prompts |
| 100 | self.name_lens = name_lens |
| 101 | self.naive_embedding = naive_embedding.cuda() |
| 102 | |
| 103 | @autocast() |
| 104 | def forward(self): |