| 80 | |
| 81 | class PromptLearner(nn.Module): |
| 82 | def __init__(self, classnames, clip_model): |
| 83 | super().__init__() |
| 84 | n_cls = len(classnames) |
| 85 | n_ctx = 16 #cfg.TRAINER.COOP.N_CTX |
| 86 | ctx_init = '' # cfg.TRAINER.COOP.CTX_INIT |
| 87 | dtype = clip_model.dtype |
| 88 | ctx_dim = clip_model.ln_final.weight.shape[0] |
| 89 | clip_imsize = clip_model.visual.input_resolution |
| 90 | cfg_imsize = 224 #cfg.INPUT.SIZE[0] |
| 91 | self.N = 4 #cfg.MODEL.N |
| 92 | assert cfg_imsize == clip_imsize, f"cfg_imsize ({cfg_imsize}) must equal to clip_imsize ({clip_imsize})" |
| 93 | |
| 94 | if ctx_init: |
| 95 | # use given words to initialize context vectors |
| 96 | ctx_init = ctx_init.replace("_", " ") |
| 97 | n_ctx = len(ctx_init.split(" ")) |
| 98 | prompt = clip.tokenize(ctx_init) |
| 99 | with torch.no_grad(): |
| 100 | embedding = clip_model.token_embedding(prompt).type(dtype) |
| 101 | ctx_vectors = embedding[0, 1 : 1 + n_ctx, :] |
| 102 | prompt_prefix = ctx_init |
| 103 | |
| 104 | else: |
| 105 | # random initialization |
| 106 | ctx_vectors = torch.empty(self.N, n_ctx, ctx_dim, dtype=dtype) |
| 107 | nn.init.normal_(ctx_vectors, std=0.02) # define the prompt to be trained |
| 108 | prompt_prefix = " ".join(["X"] * n_ctx) |
| 109 | |
| 110 | print(f'Initial context: "{prompt_prefix}"') |
| 111 | print(f"Number of context words (tokens): {n_ctx}") |
| 112 | |
| 113 | self.ctx = nn.Parameter(ctx_vectors) # to be optimized |
| 114 | |
| 115 | |
| 116 | classnames = [name.replace("_", " ") for name in classnames] |
| 117 | name_lens = [len(_tokenizer.encode(name)) for name in classnames] |
| 118 | prompts = [prompt_prefix + " " + name + "." for name in classnames] |
| 119 | |
| 120 | tokenized_prompts = torch.cat([clip.tokenize(p) for p in prompts]) |
| 121 | tokenized_prompts = tokenized_prompts.repeat(self.N,1) |
| 122 | |
| 123 | |
| 124 | with torch.no_grad(): |
| 125 | embedding = clip_model.token_embedding(tokenized_prompts).type(dtype) |
| 126 | |
| 127 | |
| 128 | # These token vectors will be saved when in save_model(), |
| 129 | # but they should be ignored in load_model() as we want to use |
| 130 | # those computed using the current class names |
| 131 | self.register_buffer("token_prefix", embedding[:, :1, :]) # SOS |
| 132 | self.register_buffer("token_suffix", embedding[:, 1 + n_ctx :, :]) # CLS, EOS |
| 133 | |
| 134 | self.n_cls = n_cls |
| 135 | self.n_ctx = n_ctx |
| 136 | self.tokenized_prompts = tokenized_prompts # torch.Tensor |
| 137 | self.name_lens = name_lens |
| 138 | self.class_token_position = 'end' #cfg.TRAINER.COOP.CLASS_TOKEN_POSITION |
| 139 | |