| 58 | |
| 59 | |
| 60 | class PromptLearner(nn.Module): |
| 61 | def __init__(self, cfg, classnames, clip_model): |
| 62 | super().__init__() |
| 63 | n_cls = len(classnames) |
| 64 | n_ctx = cfg.TRAINER.PLOT.N_CTX |
| 65 | ctx_init = cfg.TRAINER.PLOT.CTX_INIT |
| 66 | dtype = clip_model.dtype |
| 67 | ctx_dim = clip_model.ln_final.weight.shape[0] |
| 68 | clip_imsize = clip_model.visual.input_resolution |
| 69 | cfg_imsize = cfg.INPUT.SIZE[0] |
| 70 | self.N = cfg.TRAINER.PLOT.N |
| 71 | assert cfg_imsize == clip_imsize, f"cfg_imsize ({cfg_imsize}) must equal to clip_imsize ({clip_imsize})" |
| 72 | |
| 73 | if ctx_init: |
| 74 | # use given words to initialize context vectors |
| 75 | ctx_init = ctx_init.replace("_", " ") |
| 76 | n_ctx = len(ctx_init.split(" ")) |
| 77 | prompt = clip.tokenize(ctx_init) |
| 78 | with torch.no_grad(): |
| 79 | embedding = clip_model.token_embedding(prompt).type(dtype) |
| 80 | ctx_vectors = embedding[0, 1 : 1 + n_ctx, :] |
| 81 | prompt_prefix = ctx_init |
| 82 | |
| 83 | else: |
| 84 | # random initialization |
| 85 | if cfg.TRAINER.PLOT.CSC: |
| 86 | print("Initializing class-specific contexts") |
| 87 | ctx_vectors = torch.empty(n_cls, n_ctx, ctx_dim, dtype=dtype) |
| 88 | else: |
| 89 | print("Initializing a generic context") |
| 90 | ctx_vectors = torch.empty(self.N, n_ctx, ctx_dim, dtype=dtype) |
| 91 | nn.init.normal_(ctx_vectors, std=0.02) # define the prompt to be trained |
| 92 | prompt_prefix = " ".join(["X"] * n_ctx) |
| 93 | |
| 94 | print(f'Initial context: "{prompt_prefix}"') |
| 95 | print(f"Number of context words (tokens): {n_ctx}") |
| 96 | |
| 97 | self.ctx = nn.Parameter(ctx_vectors) # to be optimized |
| 98 | |
| 99 | |
| 100 | classnames = [name.replace("_", " ") for name in classnames] |
| 101 | name_lens = [len(_tokenizer.encode(name)) for name in classnames] |
| 102 | prompts = [prompt_prefix + " " + name + "." for name in classnames] |
| 103 | |
| 104 | tokenized_prompts = torch.cat([clip.tokenize(p) for p in prompts]) |
| 105 | tokenized_prompts = tokenized_prompts.repeat(self.N,1) |
| 106 | # tokenized_prompts3.view(3,100,77) |
| 107 | |
| 108 | with torch.no_grad(): |
| 109 | embedding = clip_model.token_embedding(tokenized_prompts).type(dtype) |
| 110 | |
| 111 | |
| 112 | # These token vectors will be saved when in save_model(), |
| 113 | # but they should be ignored in load_model() as we want to use |
| 114 | # those computed using the current class names |
| 115 | self.register_buffer("token_prefix", embedding[:, :1, :]) # SOS |
| 116 | self.register_buffer("token_suffix", embedding[:, 1 + n_ctx :, :]) # CLS, EOS |
| 117 | |