| 87 | |
| 88 | |
| 89 | def build_cache_model(cfg, clip_model, train_loader_cache): |
| 90 | if cfg['load_cache'] == False: |
| 91 | cache_keys = [] |
| 92 | cache_values = [] |
| 93 | cache_keymaps = [] |
| 94 | |
| 95 | with torch.no_grad(): |
| 96 | # Data augmentation for the cache model |
| 97 | for augment_idx in range(cfg['augment_epoch']): |
| 98 | train_features = [] |
| 99 | train_images_features_map = [] |
| 100 | |
| 101 | print('Augment Epoch: {:} / {:}'.format(augment_idx, cfg['augment_epoch'])) |
| 102 | for i, (images, target) in enumerate(tqdm(train_loader_cache)): |
| 103 | images = images.cuda() |
| 104 | _, img_feat, text_feat, img_feat_map = clip_model(images) |
| 105 | |
| 106 | train_features.append(img_feat) |
| 107 | train_images_features_map.append(img_feat_map) |
| 108 | if augment_idx == 0: |
| 109 | target = target.cuda() |
| 110 | cache_values.append(target) |
| 111 | if i ==0: |
| 112 | cache_plot_text=text_feat |
| 113 | |
| 114 | cache_keys.append(torch.cat(train_features, dim=0).unsqueeze(0)) |
| 115 | cache_keymaps.append(torch.cat(train_images_features_map, dim=1).unsqueeze(0)) |
| 116 | |
| 117 | cache_keys = torch.cat(cache_keys, dim=0).mean(dim=0) |
| 118 | cache_keys /= cache_keys.norm(dim=-1, keepdim=True) |
| 119 | cache_keys = cache_keys.permute(1, 0) |
| 120 | |
| 121 | cache_keymaps = torch.cat(cache_keymaps, dim=0).mean(dim=0) |
| 122 | cache_keymaps /= cache_keymaps.norm(dim=-1, keepdim=True) |
| 123 | cache_keymaps = cache_keymaps.permute(2, 0, 1) # d x M x B |
| 124 | |
| 125 | cache_values = F.one_hot(torch.cat(cache_values, dim=0)).half() |
| 126 | |
| 127 | torch.save(cache_keys, cfg['cache_dir'] + '/keys_' + str(cfg['shots']) + "shots.pt") |
| 128 | torch.save(cache_values, cfg['cache_dir'] + '/values_' + str(cfg['shots']) + "shots.pt") |
| 129 | torch.save(cache_keymaps, cfg['cache_dir'] + '/keymaps_' + str(cfg['shots']) + "shots.pt") |
| 130 | torch.save(cache_plot_text, cfg['cache_dir'] + '/plot_text_' + str(cfg['shots']) + "shots.pt") |
| 131 | |
| 132 | else: |
| 133 | cache_keys = torch.load(cfg['cache_dir'] + '/keys_' + str(cfg['shots']) + "shots.pt") |
| 134 | cache_values = torch.load(cfg['cache_dir'] + '/values_' + str(cfg['shots']) + "shots.pt") |
| 135 | cache_keymaps = torch.load(cfg['cache_dir'] + '/keymaps_' + str(cfg['shots']) + "shots.pt") |
| 136 | cache_plot_text = torch.load(cfg['cache_dir'] + '/plot_text_' + str(cfg['shots']) + "shots.pt") |
| 137 | |
| 138 | return cache_keys, cache_values, cache_keymaps, cache_plot_text |
| 139 | |
| 140 | |
| 141 | def pre_load_features(cfg, split, clip_model, loader): |