ChannelDiversification/test.py at main · rucv/ChannelDiversification · GitHub

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
import argparse
import os
import shutil
import time

import torch
import torch.nn as nn
import argparse
import os
import shutil
import time

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import torch.optim
import torch.utils.data
import torchvision.transforms as transforms
import torchvision.datasets as datasets
from torch.autograd import Variable
from ptflops import get_model_complexity_info
from network import EmbeddingNet
#from network import ResNet

parser = argparse.ArgumentParser(description='PyTorch WideResNet Training')
parser.add_argument('--dataset', default='cifar100', type=str,
                    help='dataset (cifar10 [default] or cifar100)')

parser.add_argument('-b', '--batch-size', default=128, type=int,
                    help='mini-batch size (default: 128)')
print("PyTorch Version: ", torch.__version__)
#print("Torchvision Version: ", torchvision.__version__)
parser.add_argument('--no-bottleneck', dest='bottleneck', action='store_false',
                    help='to use basicblock for CIFAR datasets (default: bottleneck)')

class ImageFolderWithPaths(datasets.CIFAR100):
     def __getitem__(self, index):
         # this is what ImageFolder normally returns
         original_tuple = super(ImageFolderWithPaths, self).__getitem__(index)
         # the image file path
         #d_list = self.test_list
         #print(d_list[index])
         path = self.data[index]
         #print(str(pathi))

         # make a new tuple that includes original and the pat
         tuple_with_path = (original_tuple + (path,))
         return tuple_with_path

global args
args = parser.parse_args()

normalize = transforms.Normalize(mean=[x/255.0 for x in [125.3, 123.0, 113.9]],
                                     std=[x/255.0 for x in [63.0, 62.1, 66.7]])


transform_train = transforms.Compose([
                transforms.ToTensor(),
        	transforms.Lambda(lambda x: F.pad(x.unsqueeze(0),
        						(4,4,4,4),mode='reflect').squeeze()),
                transforms.ToPILImage(),
                transforms.RandomCrop(32),
                transforms.RandomHorizontalFlip(),
                transforms.ToTensor(),
                normalize,
                ])

transform_test = transforms.Compose([
        transforms.ToTensor(),
        normalize
        ])

kwargs = {'num_workers': 1, 'pin_memory': True}
assert(args.dataset == 'cifar10' or args.dataset == 'cifar100')

train_loader = torch.utils.data.DataLoader(
        datasets.__dict__[args.dataset.upper()]('../data', train=True, download=True,
                         transform=transform_train),
        batch_size=args.batch_size, shuffle=True, **kwargs)

val_loader = torch.utils.data.DataLoader(
        ImageFolderWithPaths('../data', train=False, transform=transform_test),
        batch_size=args.batch_size, shuffle=False, **kwargs)


def test_model(model,val_loader):

    best_model_wts = torch.load("./runs/ResNet-28-10/model_best.pth.tar")
    #checkpoint = torch.load('model_best.pth.tar')
    model = torch.nn.DataParallel(model).cuda()
    model.load_state_dict(best_model_wts['state_dict'])
    model.eval()
    #model = torch.nn.DataParallel(model).cuda()
    for param_tensor in model.state_dict():
       print(param_tensor, "\t", model.state_dict()[param_tensor].size())

    total = 0
    correct = 0

    with torch.no_grad():
        for inputs, labels, path in val_loader:
            inputs = inputs.to(device)
            labels = labels.to(device)
            outputs = model.forward(inputs)
            _,preds = torch.max(outputs,1)
            c = (preds == labels).squeeze()

            total += labels.size(0)
            correct += (preds == labels).sum().item()
            path_list = list(path)
            l = 0


    print('Accuracy of the network on the 10000 test images: %d %%' % (
    100 * correct / total))


model_ft = ResNet(args.bottleneck)


device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

model_ft = model_ft.to(device)

test_model(model_ft, val_loader)