-
Notifications
You must be signed in to change notification settings - Fork 2
/
test_gvcnn.py
95 lines (74 loc) · 2.78 KB
/
test_gvcnn.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
# -*- coding: utf-8 -*-
from __future__ import print_function
import datasets
import os
import sys
sys.path.append('datasets')
sys.path.append('models')
sys.path.append('utils')
from models import gvcnn
from datasets import modelnet40
from utils import meter
import utils.config
import torch
import torch.utils.data
from torch.autograd import Variable
def test(test_loader, model, cfg):
"""
test for one epoch on the testing set
"""
batch_time = meter.TimeMeter(True)
data_time = meter.TimeMeter(True)
prec = meter.ClassErrorMeter(topk=[1], accuracy=True)
mAP = meter.mAPMeter()
# training mode
model.eval()
for i, (shapes, labels) in enumerate(test_loader):
batch_time.reset()
# bz x 12 x 3 x 224 x 224
labels = labels.long().view(-1)
shapes = Variable(shapes)
labels = Variable(labels)
# shift data to GPU
if cfg.cuda:
shapes = shapes.cuda()
labels = labels.cuda()
# forward, backward optimize
preds = model(shapes)
labels_oh = torch.zeros(labels.data.size(0), cfg.class_num)\
.scatter_(1, labels.cpu().data.unsqueeze(1), 1)
mAP.add(preds.data, labels_oh)
prec.add(preds.data, labels.data)
if i % cfg.print_freq == 0:
print('[{0}/{1}]\t'
'Batch Time {batch_time:.3f}\t'
'Epoch Time {data_time:.3f}\t'
'Prec@1 {top1:.3f}\t'.format(
i, len(test_loader), batch_time=batch_time.value(),
data_time=data_time.value(), top1=prec.value(1)))
print('mean class accuracy : {0} '.format(prec.value(1)))
print('mAP: %f' % mAP.value())
def main():
cfg = utils.config.config()
os.environ['CUDA_VISIBLE_DEVICES'] = cfg.gpu_id
test_dataset = datasets.modelnet40.Modelnet40_dataset(cfg, status='test')
print('number of test samples is: ', len(test_dataset))
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=cfg.batch_size,
shuffle=False, num_workers=cfg.workers)
# create model
model = gvcnn.GVCNN(pretrained=True, aux_logits=cfg.have_aux,
transform_input=False, num_classes=40,
n_views=cfg.data_views, with_group=cfg.with_group)
print('loading model from {0}'.format(cfg.ckpt_model))
checkpoint = torch.load(cfg.ckpt_model)
model.load_state_dict(checkpoint['model_param_best'])
for p in model.parameters():
p.requires_grad = False # to avoid computation
print('GVCNN: ')
print(model)
if cfg.cuda:
print('shift model and criterion to GPU .. ')
model = model.cuda()
test(test_loader, model, cfg)
if __name__ == '__main__':
main()