| 55 | |
| 56 | class MultiModalPromptLearner(nn.Module): |
| 57 | def __init__(self, cfg, classnames, clip_model): |
| 58 | super().__init__() |
| 59 | n_cls = len(classnames) |
| 60 | n_ctx = cfg.TRAINER.PLOTPP.N_CTX |
| 61 | n_ctx_vision = cfg.TRAINER.PLOTPP.N_CTX_V # the number of vision context tokens |
| 62 | ctx_init_flag = cfg.TRAINER.PLOTPP.CTX_INIT |
| 63 | dtype = clip_model.dtype |
| 64 | ctx_dim = clip_model.ln_final.weight.shape[0] |
| 65 | clip_imsize = clip_model.visual.input_resolution |
| 66 | cfg_imsize = cfg.INPUT.SIZE[0] |
| 67 | M = cfg.TRAINER.PLOTPP.M #the number of our visual prompts |
| 68 | N = cfg.TRAINER.PLOTPP.N # the number of our text prompts |
| 69 | self.M = M |
| 70 | self.N = N |
| 71 | assert cfg_imsize == clip_imsize, f"cfg_imsize ({cfg_imsize}) must equal to clip_imsize ({clip_imsize})" |
| 72 | |
| 73 | template_dict = {'Caltech101': ["a photo of a","this is a photo","this is picture of","one picture of a"], |
| 74 | 'DescribableTextures':['a photo of a texture', "this is a photo texture","this is a picture texture","one picture of a texture"], |
| 75 | 'EuroSAT':['a centered satellite photo of', 'a centered satellite picture of','this is centered satellite photo of','one centered satellite photo of a'], |
| 76 | 'FGVCAircraft':['a photo of an aircraft','a picture of an aircraft','this is aircraft picture of','one picture of an aircraft'], |
| 77 | 'Food101':['a photo of a food', 'this is a food photo', ' this is food picture of','one picture of a food'], |
| 78 | 'ImageNet':["a photo of a","this is a photo ","this is a","one picture of a"], |
| 79 | 'OxfordFlowers':['a photo of a flower', 'one picture of a flower','this is flower picture of','one picture of a flower'], |
| 80 | 'OxfordPets':['a photo of a pet', 'one picture of a pet','this is pet picture of','one picture of a pet'], |
| 81 | 'StanfordCars':["a photo of a","this is a photo ","this is picture of","one picture of a"], |
| 82 | 'SUN397':["a photo of a","this is a photo","this is picture of","one picture of a"], |
| 83 | 'UCF101':['a photo of a person doing', 'this is a photo people doing', 'this is picture of people doing', 'one picture of a person doing'],} |
| 84 | |
| 85 | if ctx_init_flag: |
| 86 | ctx_list = template_dict[cfg.DATASET.NAME] |
| 87 | n_ctx = len(ctx_list[0].split()) |
| 88 | ctx_vectors_list = [] |
| 89 | prompt_prefix_list = [] |
| 90 | |
| 91 | for i in range(N): |
| 92 | ctx_init = ctx_list[i].replace("_", " ") |
| 93 | prompt = clip.tokenize(ctx_init) |
| 94 | with torch.no_grad(): |
| 95 | embedding = clip_model.token_embedding(prompt).type(dtype) |
| 96 | ctx_vectors = embedding[0, 1 : 1 + n_ctx, :] |
| 97 | ctx_vectors_list.append(ctx_vectors) |
| 98 | prompt_prefix = ctx_init |
| 99 | prompt_prefix_list.append(prompt_prefix) |
| 100 | ctx_vision_vectors = torch.empty(M, n_ctx_vision ,768, dtype=dtype) |
| 101 | nn.init.normal_(ctx_vision_vectors, std=0.02) |
| 102 | ctx_vectors = torch.stack(ctx_vectors_list) |
| 103 | |
| 104 | else: |
| 105 | ctx_vectors = torch.empty(N, n_ctx, ctx_dim, dtype=dtype) |
| 106 | ctx_vision_vectors = torch.empty(M, n_ctx_vision ,768, dtype=dtype) |
| 107 | nn.init.normal_(ctx_vectors, std=0.02) |
| 108 | nn.init.normal_(ctx_vision_vectors, std=0.02) |
| 109 | prompt_prefix = " ".join(["X"] * n_ctx) |
| 110 | |
| 111 | self.ctx = nn.Parameter(ctx_vectors) # parameters of text prompt to be learned |
| 112 | self.ctx_vision = nn.Parameter(ctx_vision_vectors) # parameters of vision prompt to be learned |
| 113 | |
| 114 | classnames = [name.replace("_", " ") for name in classnames] |