-
Notifications
You must be signed in to change notification settings - Fork 41
/
main_imagenet.py
164 lines (116 loc) · 6.05 KB
/
main_imagenet.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
import os
import random
import argparse
import yaml
from tqdm import tqdm
import torch
import torch.nn.functional as F
import torch.nn as nn
from datasets.imagenet import ImageNet
import clip
from utils import *
def get_arguments():
parser = argparse.ArgumentParser()
parser.add_argument('--config', dest='config', help='settings of Tip-Adapter in yaml format')
args = parser.parse_args()
return args
def run_tip_adapter(cfg, cache_keys, cache_values, test_features, test_labels, clip_weights):
# Zero-shot CLIP
clip_logits = 100. * test_features @ clip_weights
acc = cls_acc(clip_logits, test_labels)
print("\n**** Zero-shot CLIP's test accuracy: {:.2f}. ****\n".format(acc))
# Tip-Adapter
beta, alpha = cfg['init_beta'], cfg['init_alpha']
affinity = test_features @ cache_keys
cache_logits = ((-1) * (beta - beta * affinity)).exp() @ cache_values
tip_logits = clip_logits + cache_logits * alpha
acc = cls_acc(tip_logits, test_labels)
print("**** Tip-Adapter's test accuracy: {:.2f}. ****\n".format(acc))
# Search Hyperparameters
_ = search_hp(cfg, cache_keys, cache_values, test_features, test_labels, clip_weights)
def run_tip_adapter_F(cfg, cache_keys, cache_values, test_features, test_labels, clip_weights, clip_model, train_loader_F):
# Enable the cached keys to be learnable
adapter = nn.Linear(cache_keys.shape[0], cache_keys.shape[1], bias=False).to(clip_model.dtype).cuda()
adapter.weight = nn.Parameter(cache_keys.t())
optimizer = torch.optim.AdamW(adapter.parameters(), lr=cfg['lr'], eps=1e-4)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, cfg['train_epoch'] * len(train_loader_F))
beta, alpha = cfg['init_beta'], cfg['init_alpha']
best_acc, best_epoch = 0.0, 0
for train_idx in range(cfg['train_epoch']):
# Train
adapter.train()
correct_samples, all_samples = 0, 0
loss_list = []
print('Train Epoch: {:} / {:}'.format(train_idx, cfg['train_epoch']))
for i, (images, target) in enumerate(tqdm(train_loader_F)):
images, target = images.cuda(), target.cuda()
with torch.no_grad():
image_features = clip_model.encode_image(images)
image_features /= image_features.norm(dim=-1, keepdim=True)
affinity = adapter(image_features)
cache_logits = ((-1) * (beta - beta * affinity)).exp() @ cache_values
clip_logits = 100. * image_features @ clip_weights
tip_logits = clip_logits + cache_logits * alpha
loss = F.cross_entropy(tip_logits, target)
acc = cls_acc(tip_logits, target)
correct_samples += acc / 100 * len(tip_logits)
all_samples += len(tip_logits)
loss_list.append(loss.item())
optimizer.zero_grad()
loss.backward()
optimizer.step()
scheduler.step()
current_lr = scheduler.get_last_lr()[0]
print('LR: {:.6f}, Acc: {:.4f} ({:}/{:}), Loss: {:.4f}'.format(current_lr, correct_samples / all_samples, correct_samples, all_samples, sum(loss_list)/len(loss_list)))
# Eval
adapter.eval()
affinity = adapter(test_features)
cache_logits = ((-1) * (beta - beta * affinity)).exp() @ cache_values
clip_logits = 100. * test_features @ clip_weights
tip_logits = clip_logits + cache_logits * alpha
acc = cls_acc(tip_logits, test_labels)
print("**** Tip-Adapter-F's test accuracy: {:.2f}. ****\n".format(acc))
if acc > best_acc:
best_acc = acc
best_epoch = train_idx
torch.save(adapter.weight, cfg['cache_dir'] + "/best_F_" + str(cfg['shots']) + "shots.pt")
adapter.weight = torch.load(cfg['cache_dir'] + "/best_F_" + str(cfg['shots']) + "shots.pt")
print(f"**** After fine-tuning, Tip-Adapter-F's best test accuracy: {best_acc:.2f}, at epoch: {best_epoch}. ****\n")
# Search Hyperparameters
_ = search_hp(cfg, affinity, cache_values, test_features, test_labels, clip_weights, adapter=adapter)
def main():
# Load config file
args = get_arguments()
assert (os.path.exists(args.config))
cfg = yaml.load(open(args.config, 'r'), Loader=yaml.Loader)
cache_dir = os.path.join('./caches', cfg['dataset'])
os.makedirs(cache_dir, exist_ok=True)
cfg['cache_dir'] = cache_dir
print("\nRunning configs.")
print(cfg, "\n")
# CLIP
clip_model, preprocess = clip.load(cfg['backbone'])
clip_model.eval()
# ImageNet dataset
random.seed(1)
torch.manual_seed(1)
print("Preparing ImageNet dataset.")
imagenet = ImageNet(cfg['root_path'], cfg['shots'], preprocess)
test_loader = torch.utils.data.DataLoader(imagenet.test, batch_size=64, num_workers=8, shuffle=False)
train_loader_cache = torch.utils.data.DataLoader(imagenet.train, batch_size=256, num_workers=8, shuffle=False)
train_loader_F = torch.utils.data.DataLoader(imagenet.train, batch_size=256, num_workers=8, shuffle=True)
# Textual features
print("Getting textual features as CLIP's classifier.")
clip_weights = clip_classifier(imagenet.classnames, imagenet.template, clip_model)
# Construct the cache model by few-shot training set
print("\nConstructing cache model by few-shot visual features and labels.")
cache_keys, cache_values = build_cache_model(cfg, clip_model, train_loader_cache)
# Pre-load test features
print("\nLoading visual features and labels from test set.")
test_features, test_labels = pre_load_features(cfg, "test", clip_model, test_loader)
# ------------------------------------------ Tip-Adapter ------------------------------------------
run_tip_adapter(cfg, cache_keys, cache_values, test_features, test_labels, clip_weights)
# ------------------------------------------ Tip-Adapter-F ------------------------------------------
run_tip_adapter_F(cfg, cache_keys, cache_values, test_features, test_labels, clip_weights, clip_model, train_loader_F)
if __name__ == '__main__':
main()