import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
import torch.backends.cudnn as cudnn

import torchvision
import torchvision.transforms as transforms
import torchvision.datasets as datasets

import os
import argparse

from torch.autograd import Variable
from extensions.utils import progress_bar
from extensions.model_refinery_wrapper import ModelRefineryWrapper
from extensions.refinery_loss import RefineryLoss

from models import ShuffleNetv2_wrapper
from models import DiracDeltaNet_wrapper


parser = argparse.ArgumentParser(description='PyTorch imagenet inference')
parser.add_argument('--datadir', help='path to dataset')
parser.add_argument('--inputdir', help='path to input model')
args = parser.parse_args()


# Data
print('==> Preparing data..')
# Data loading code
valdir = os.path.join(args.datadir, 'val')
    
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                 std=[0.229, 0.224, 0.225])


transform_test = transforms.Compose([
            transforms.Resize(256),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            normalize,
        ])

#imagenet
testset = datasets.ImageFolder(valdir, transform_test)
num_classes=1000

testloader = torch.utils.data.DataLoader(testset, batch_size=1000, shuffle=False, pin_memory=True, num_workers=30)


use_cuda = torch.cuda.is_available()

print('Using input path: %s' % args.inputdir)
checkpoint = torch.load(args.inputdir)
init_net = checkpoint['net']
net=init_net.to('cpu')

label_refinery=torch.load('./resnet50.t7')
net = ModelRefineryWrapper(net, label_refinery)

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

if torch.cuda.device_count() > 1:
    print("Let's use", torch.cuda.device_count(), "GPUs!")
    net = nn.DataParallel(net)

net=net.to(device)

criterion = RefineryLoss()



def accuracy(output, target, topk=(1,)):
    """Computes the precision@k for the specified values of k"""
    with torch.no_grad():
        maxk = max(topk)
        batch_size = target.size(0)

        _, pred = output.topk(maxk, 1, True, True)
        pred = pred.t()
        correct = pred.eq(target.view(1, -1).expand_as(pred))

        res = []
        for k in topk:
            correct_k = correct[:k].view(-1).float().sum(0, keepdim=True)
            res.append(correct_k)
        return res


def test():
    net.eval()
    criterion.eval()
    test_loss = 0
    correct_1 = 0
    correct_5 = 0
    total = 0
    for batch_idx, (inputs, targets) in enumerate(testloader):
        if use_cuda:
            inputs, targets = inputs.cuda(device), targets.cuda(device)
        with torch.no_grad():
            outputs = net(inputs)
        loss = criterion(outputs, targets)

        if isinstance(loss, tuple):
            loss_value, outputs = loss
        else:
            loss_value = loss

        test_loss += loss_value.item()
        prec1, prec5 = accuracy(outputs, targets, topk=(1, 5))
        total += targets.size(0)
        correct_1 += prec1
        correct_5 += prec5

        progress_bar(batch_idx, len(testloader), 'Loss: %.3f | Acc: %.3f%% (%d/%d)'
            % (test_loss/(batch_idx+1), 100.*float(correct_1)/float(total), correct_1, total))

    return 100.*float(correct_1)/float(total),100.*float(correct_5)/float(total),test_loss

acc1,acc5,loss=test()
print('top-1 accuracy: {0:.3f}%, top-5 accuracy: {1:.3f}%'.format(acc1,acc5))

[((23, 9, 23, 74), 'argparse.ArgumentParser', 'argparse.ArgumentParser', (), '', False, 'import argparse\n'), ((32, 9, 32, 42), 'os.path.join', 'os.path.join', ({(32, 22, 32, 34): 'args.datadir', (32, 36, 32, 41): '"""val"""'}, {}), "(args.datadir, 'val')", False, 'import os\n'), ((34, 12, 35, 59), 'torchvision.transforms.Normalize', 'transforms.Normalize', (), '', True, 'import torchvision.transforms as transforms\n'), ((46, 10, 46, 54), 'torchvision.datasets.ImageFolder', 'datasets.ImageFolder', ({(46, 31, 46, 37): 'valdir', (46, 39, 46, 53): 'transform_test'}, {}), '(valdir, transform_test)', True, 'import torchvision.datasets as datasets\n'), ((49, 13, 49, 114), 'torch.utils.data.DataLoader', 'torch.utils.data.DataLoader', (), '', False, 'import torch\n'), ((52, 11, 52, 36), 'torch.cuda.is_available', 'torch.cuda.is_available', ({}, {}), '()', False, 'import torch\n'), ((55, 13, 55, 38), 'torch.load', 'torch.load', ({(55, 24, 55, 37): 'args.inputdir'}, {}), '(args.inputdir)', False, 'import torch\n'), ((59, 15, 59, 42), 'torch.load', 'torch.load', ({(59, 26, 59, 41): '"""./resnet50.t7"""'}, {}), "('./resnet50.t7')", False, 'import torch\n'), ((60, 6, 60, 47), 'extensions.model_refinery_wrapper.ModelRefineryWrapper', 'ModelRefineryWrapper', ({(60, 27, 60, 30): 'net', (60, 32, 60, 46): 'label_refinery'}, {}), '(net, label_refinery)', False, 'from extensions.model_refinery_wrapper import ModelRefineryWrapper\n'), ((71, 12, 71, 26), 'extensions.refinery_loss.RefineryLoss', 'RefineryLoss', ({}, {}), '()', False, 'from extensions.refinery_loss import RefineryLoss\n'), ((65, 3, 65, 28), 'torch.cuda.device_count', 'torch.cuda.device_count', ({}, {}), '()', False, 'import torch\n'), ((67, 10, 67, 30), 'torch.nn.DataParallel', 'nn.DataParallel', ({(67, 26, 67, 29): 'net'}, {}), '(net)', True, 'import torch.nn as nn\n'), ((39, 12, 39, 34), 'torchvision.transforms.Resize', 'transforms.Resize', ({(39, 30, 39, 33): '256'}, {}), '(256)', True, 'import torchvision.transforms as transforms\n'), ((40, 12, 40, 38), 'torchvision.transforms.CenterCrop', 'transforms.CenterCrop', ({(40, 34, 40, 37): '224'}, {}), '(224)', True, 'import torchvision.transforms as transforms\n'), ((41, 12, 41, 33), 'torchvision.transforms.ToTensor', 'transforms.ToTensor', ({}, {}), '()', True, 'import torchvision.transforms as transforms\n'), ((62, 34, 62, 59), 'torch.cuda.is_available', 'torch.cuda.is_available', ({}, {}), '()', False, 'import torch\n'), ((66, 23, 66, 48), 'torch.cuda.device_count', 'torch.cuda.device_count', ({}, {}), '()', False, 'import torch\n'), ((77, 9, 77, 24), 'torch.no_grad', 'torch.no_grad', ({}, {}), '()', False, 'import torch\n'), ((102, 13, 102, 28), 'torch.no_grad', 'torch.no_grad', ({}, {}), '()', False, 'import torch\n')]