-
Notifications
You must be signed in to change notification settings - Fork 5
/
evaluate_imagenet_models.py
75 lines (56 loc) · 2.52 KB
/
evaluate_imagenet_models.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
import torch
import argparse
from utils.models import load_model
import os
from tqdm import tqdm
from utils.data_loaders_imgnet_id import ImageNetDLoader
def get_args_parser():
parser = argparse.ArgumentParser('Evaluate vision models trained on ImageNet', add_help=False)
parser.add_argument('--model', default='vgg19', type=str, metavar='MODEL', help='Name of model to test')
return parser
def test_model(test_loader_dict, model, args):
model.eval()
for datatype, test_loader in test_loader_dict.items():
top1_cnt = 0
top5_cnt = 0
cnt = 0
result_dir = f"result/{datatype}/"
os.makedirs(result_dir, exist_ok=True)
pbar_prefix = 'Testing on ' + str(datatype)
with torch.no_grad():
for sample, example in tqdm(enumerate(test_loader), total=len(test_loader), desc=pbar_prefix):
images = example['images']
target = example['labels']
images = images.cuda()
output = model(images)
maxk = 5
batch_size = target[0].shape[0]
_, pred = output.topk(maxk, 1, True, True)
pred = pred.t()
for k in range(len(target)):
correct = pred.eq(target[k].cuda().view(1, -1).expand_as(pred))
top1_cnt += correct[:1].view(-1).float().sum(0, keepdim=True)
top5_cnt += correct[:5].reshape(-1).float().sum(0, keepdim=True)
cnt += batch_size
print("Test accuracy: %d images top 1: %.5f, top 5: %.5f" %
(cnt, top1_cnt * 100.0 / cnt, top5_cnt * 100.0 / cnt))
with open(f"{result_dir}/imagenet_model_accuracy.txt", 'a+') as txt_file:
txt_file.write(f'{args.model}\t{cnt}\t{top1_cnt.item() * 100.0 / cnt}\t{top5_cnt.item() * 100.0 / cnt}\n')
def test():
args = get_args_parser()
args = args.parse_args()
imagenet_d_dir = './data/ImageNet-D/'
model, test_transform = load_model(args.model)
model.cuda()
model.eval()
test_loader_all = dict()
subset_list = ['background', 'texture', 'material']
for subset in subset_list:
test_loader = torch.utils.data.DataLoader(
ImageNetDLoader(os.path.join(imagenet_d_dir, subset), transform=test_transform),
batch_size=100, shuffle=True,
num_workers=10, pin_memory=True)
test_loader_all[subset] = test_loader
test_model(test_loader_all, model, args)
if __name__ == "__main__":
test()