import argparse
import os
#os.environ['CUDA_VISIBLE_DEVICES']='2'
import shutil
import time

import torch
import torch.nn as nn
import torch.nn.parallel
#import torch.backends.cudnn as cudnn
import torch.optim
import torch.utils.data
import torchvision
import torchvision.transforms as transforms
import torchvision.datasets as datasets
import torchvision.models as models
import matplotlib.pyplot as plt
import numpy as np
from MobileNetV2 import MobileNetV2
import pdb


model_names = sorted(name for name in models.__dict__
    if name.islower() and not name.startswith("__")
    and callable(models.__dict__[name]))
model_names.append('mobilenetv2')


parser = argparse.ArgumentParser(description='PyTorch ImageNet Training')
parser.add_argument('--data', metavar='DIR',
                    help='path to dataset')
parser.add_argument('--arch', '-a', metavar='ARCH', default='mobilenetv2',
                    choices=model_names,
                    help='model architecture: ' +
                        ' | '.join(model_names) +
                        ' (default: resnet18)')
parser.add_argument('-j', '--workers', default=4, type=int, metavar='N',
                    help='number of data loading workers (default: 4)')
parser.add_argument('--epochs', default=90, type=int, metavar='N',
                    help='number of total epochs to run')
parser.add_argument('--start-epoch', default=0, type=int, metavar='N',
                    help='manual epoch number (useful on restarts)')
parser.add_argument('-b', '--batch-size', default=4, type=int,
                    metavar='N',help='input batch size for training (default: 20')
parser.add_argument('--lr', '--learning-rate', default=0.1, type=float,
                    metavar='LR', help='initial learning rate')
parser.add_argument('--momentum', default=0.9, type=float, metavar='M',
                    help='momentum')
parser.add_argument('--weight-decay', '--wd', default=1e-4, type=float,
                    metavar='W', help='weight decay (default: 1e-4)')
parser.add_argument('--print-freq', '-p', default=10, type=int,
                    metavar='N', help='print frequency (default: 10)')
parser.add_argument('--resume', default='', type=str, metavar='PATH',
                    help='path to latest checkpoint (default: none)')
parser.add_argument('--evaluate', default='', type=str, metavar='PATH',
                    help='path to evaluate model (default: none)')
parser.add_argument('--pretrained', dest='pretrained', action='store_true',
                    help='use pre-trained model')

best_prec1 = 0


MEAN_COEF = [0.48502, 0.45796, 0.40760]
DIV_COEF = [0.23068, 0.23068, 0.23068]

transform=transforms.Compose([transforms.ToTensor(),transforms.Normalize(mean=MEAN_COEF,std=DIV_COEF)])
train_data=datasets.CIFAR10('data',train=True,download=True,transform=transform)
test_data=datasets.CIFAR10('data',train=False,download=True,transform=transform)

# functions to show an image

def imshow(img):
    npimg = img * np.array(DIV_COEF).reshape([3,1,1]) \
          + np.array(MEAN_COEF).reshape([3,1,1])            # Un-normalize
    plt.imshow(np.transpose(npimg, (1, 2, 0)))


def main():
    global args, best_prec1
    args = parser.parse_args()

    # create model
    if args.pretrained:
        print("=> using pre-trained model '{}'".format(args.arch))
        model = models.__dict__[args.arch](pretrained=True)
    else:
        print("=> creating model '{}'".format(args.arch))
        if args.arch.startswith('mobilenetv1'):
            model = MobileNetV1()
        elif args.arch.startswith('mobilenetv2'):
            model = MobileNetV2()
        else:
            model = models.__dict__[args.arch]()
        print(model)

    if args.arch.startswith('alexnet') or args.arch.startswith('vgg'):
        model.features = torch.nn.DataParallel(model.features)
        model
    else:
         model=model

    # define loss function (criterion) and optimizer
    criterion = nn.CrossEntropyLoss()

    optimizer = torch.optim.SGD(model.parameters(), args.lr,
                                momentum=args.momentum,
                                weight_decay=args.weight_decay)


    if args.evaluate:
        if os.path.isfile(args.evaluate):
            print("=> loading model '{}'".format(args.evaluate))
            model.load_state_dict(torch.load(args.evaluate))
        else:
            print("=> no model found at '{}'".format(args.evaluate))


    #cudnn.benchmark = True

    val_loader = torch.utils.data.DataLoader(train_data,
        batch_size=args.batch_size,
        shuffle=False,
        num_workers=args.workers, pin_memory=True)
    

    if args.evaluate:
        validate(val_loader, model, criterion)
        return


def validate(val_loader, model, criterion):
    batch_time = AverageMeter()
    losses = AverageMeter()
    top1 = AverageMeter()
    top5 = AverageMeter()

    # switch to evaluate mode
    model.eval()

    end = time.time()
    for i, (inputs, target) in enumerate(val_loader):
        #pdb.set_trace()
        permute = [2, 1, 0]
        inputs = inputs[:,permute]

        #imshow(torchvision.utils.make_grid(inputs))
        #plt.show()


        target = target
        inputs = inputs
        input_var = torch.autograd.Variable(inputs, volatile=True)
        target_var = torch.autograd.Variable(target, volatile=True)
        
        # compute output
        output = model(input_var)

        
        #pdb.set_trace()
        loss = criterion(output.squeeze(), target_var.long())

        # measure accuracy and record loss
        prec1, prec5 = accuracy(output.squeeze().data, target, topk=(1, 5))
        losses.update(loss.data[0], inputs.size(0))
        top1.update(prec1[0], inputs.size(0))
        top5.update(prec5[0], inputs.size(0))

        # measure elapsed time
        batch_time.update(time.time() - end)
        end = time.time()

        if i % args.print_freq == 0:
            print('Test: [{0}/{1}]\t'
                  'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
                  'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
                  'Prec@1 {top1.val:.3f} ({top1.avg:.3f})\t'
                  'Prec@5 {top5.val:.3f} ({top5.avg:.3f})'.format(
                   i, len(val_loader), batch_time=batch_time, loss=losses,
                   top1=top1, top5=top5))

    print(' * Prec@1 {top1.avg:.3f} Prec@5 {top5.avg:.3f}'
          .format(top1=top1, top5=top5))



    
    return top1.avg


class AverageMeter(object):
    """Computes and stores the average and current value"""
    def __init__(self):
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count


def accuracy(output, target, topk=(1,)):
    """Computes the precision@k for the specified values of k"""
    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)
        res.append(correct_k.mul_(100.0 / batch_size))
    return res


if __name__ == '__main__':
    main()
eb2a7cc3fbd22d7c0560dca53a63a1074eb547d8
<mask token>


class AverageMeter(object):
    <mask token>

    def __init__(self):
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count


<mask token>

<mask token>


def main():
    global args, best_prec1
    args = parser.parse_args()
    if args.pretrained:
        print("=> using pre-trained model '{}'".format(args.arch))
        model = models.__dict__[args.arch](pretrained=True)
    else:
        print("=> creating model '{}'".format(args.arch))
        if args.arch.startswith('mobilenetv1'):
            model = MobileNetV1()
        elif args.arch.startswith('mobilenetv2'):
            model = MobileNetV2()
        else:
            model = models.__dict__[args.arch]()
        print(model)
    if args.arch.startswith('alexnet') or args.arch.startswith('vgg'):
        model.features = torch.nn.DataParallel(model.features)
        model
    else:
        model = model
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.SGD(model.parameters(), args.lr, momentum=args.
        momentum, weight_decay=args.weight_decay)
    if args.evaluate:
        if os.path.isfile(args.evaluate):
            print("=> loading model '{}'".format(args.evaluate))
            model.load_state_dict(torch.load(args.evaluate))
        else:
            print("=> no model found at '{}'".format(args.evaluate))
    val_loader = torch.utils.data.DataLoader(train_data, batch_size=args.
        batch_size, shuffle=False, num_workers=args.workers, pin_memory=True)
    if args.evaluate:
        validate(val_loader, model, criterion)
        return


def validate(val_loader, model, criterion):
    batch_time = AverageMeter()
    losses = AverageMeter()
    top1 = AverageMeter()
    top5 = AverageMeter()
    model.eval()
    end = time.time()
    for i, (inputs, target) in enumerate(val_loader):
        permute = [2, 1, 0]
        inputs = inputs[:, permute]
        target = target
        inputs = inputs
        input_var = torch.autograd.Variable(inputs, volatile=True)
        target_var = torch.autograd.Variable(target, volatile=True)
        output = model(input_var)
        loss = criterion(output.squeeze(), target_var.long())
        prec1, prec5 = accuracy(output.squeeze().data, target, topk=(1, 5))
        losses.update(loss.data[0], inputs.size(0))
        top1.update(prec1[0], inputs.size(0))
        top5.update(prec5[0], inputs.size(0))
        batch_time.update(time.time() - end)
        end = time.time()
        if i % args.print_freq == 0:
            print(
                'Test: [{0}/{1}]\tTime {batch_time.val:.3f} ({batch_time.avg:.3f})\tLoss {loss.val:.4f} ({loss.avg:.4f})\tPrec@1 {top1.val:.3f} ({top1.avg:.3f})\tPrec@5 {top5.val:.3f} ({top5.avg:.3f})'
                .format(i, len(val_loader), batch_time=batch_time, loss=
                losses, top1=top1, top5=top5))
    print(' * Prec@1 {top1.avg:.3f} Prec@5 {top5.avg:.3f}'.format(top1=top1,
        top5=top5))
    return top1.avg


class AverageMeter(object):
    """Computes and stores the average and current value"""

    def __init__(self):
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count


def accuracy(output, target, topk=(1,)):
    """Computes the precision@k for the specified values of k"""
    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)
        res.append(correct_k.mul_(100.0 / batch_size))
    return res


<mask token>

<mask token>
model_names.append('mobilenetv2')
<mask token>
parser.add_argument('--data', metavar='DIR', help='path to dataset')
parser.add_argument('--arch', '-a', metavar='ARCH', default='mobilenetv2',
    choices=model_names, help='model architecture: ' + ' | '.join(
    model_names) + ' (default: resnet18)')
parser.add_argument('-j', '--workers', default=4, type=int, metavar='N',
    help='number of data loading workers (default: 4)')
parser.add_argument('--epochs', default=90, type=int, metavar='N', help=
    'number of total epochs to run')
parser.add_argument('--start-epoch', default=0, type=int, metavar='N', help
    ='manual epoch number (useful on restarts)')
parser.add_argument('-b', '--batch-size', default=4, type=int, metavar='N',
    help='input batch size for training (default: 20')
parser.add_argument('--lr', '--learning-rate', default=0.1, type=float,
    metavar='LR', help='initial learning rate')
parser.add_argument('--momentum', default=0.9, type=float, metavar='M',
    help='momentum')
parser.add_argument('--weight-decay', '--wd', default=0.0001, type=float,
    metavar='W', help='weight decay (default: 1e-4)')
parser.add_argument('--print-freq', '-p', default=10, type=int, metavar='N',
    help='print frequency (default: 10)')
parser.add_argument('--resume', default='', type=str, metavar='PATH', help=
    'path to latest checkpoint (default: none)')
parser.add_argument('--evaluate', default='', type=str, metavar='PATH',
    help='path to evaluate model (default: none)')
parser.add_argument('--pretrained', dest='pretrained', action='store_true',
    help='use pre-trained model')
<mask token>


def imshow(img):
    npimg = img * np.array(DIV_COEF).reshape([3, 1, 1]) + np.array(MEAN_COEF
        ).reshape([3, 1, 1])
    plt.imshow(np.transpose(npimg, (1, 2, 0)))


def main():
    global args, best_prec1
    args = parser.parse_args()
    if args.pretrained:
        print("=> using pre-trained model '{}'".format(args.arch))
        model = models.__dict__[args.arch](pretrained=True)
    else:
        print("=> creating model '{}'".format(args.arch))
        if args.arch.startswith('mobilenetv1'):
            model = MobileNetV1()
        elif args.arch.startswith('mobilenetv2'):
            model = MobileNetV2()
        else:
            model = models.__dict__[args.arch]()
        print(model)
    if args.arch.startswith('alexnet') or args.arch.startswith('vgg'):
        model.features = torch.nn.DataParallel(model.features)
        model
    else:
        model = model
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.SGD(model.parameters(), args.lr, momentum=args.
        momentum, weight_decay=args.weight_decay)
    if args.evaluate:
        if os.path.isfile(args.evaluate):
            print("=> loading model '{}'".format(args.evaluate))
            model.load_state_dict(torch.load(args.evaluate))
        else:
            print("=> no model found at '{}'".format(args.evaluate))
    val_loader = torch.utils.data.DataLoader(train_data, batch_size=args.
        batch_size, shuffle=False, num_workers=args.workers, pin_memory=True)
    if args.evaluate:
        validate(val_loader, model, criterion)
        return


def validate(val_loader, model, criterion):
    batch_time = AverageMeter()
    losses = AverageMeter()
    top1 = AverageMeter()
    top5 = AverageMeter()
    model.eval()
    end = time.time()
    for i, (inputs, target) in enumerate(val_loader):
        permute = [2, 1, 0]
        inputs = inputs[:, permute]
        target = target
        inputs = inputs
        input_var = torch.autograd.Variable(inputs, volatile=True)
        target_var = torch.autograd.Variable(target, volatile=True)
        output = model(input_var)
        loss = criterion(output.squeeze(), target_var.long())
        prec1, prec5 = accuracy(output.squeeze().data, target, topk=(1, 5))
        losses.update(loss.data[0], inputs.size(0))
        top1.update(prec1[0], inputs.size(0))
        top5.update(prec5[0], inputs.size(0))
        batch_time.update(time.time() - end)
        end = time.time()
        if i % args.print_freq == 0:
            print(
                'Test: [{0}/{1}]\tTime {batch_time.val:.3f} ({batch_time.avg:.3f})\tLoss {loss.val:.4f} ({loss.avg:.4f})\tPrec@1 {top1.val:.3f} ({top1.avg:.3f})\tPrec@5 {top5.val:.3f} ({top5.avg:.3f})'
                .format(i, len(val_loader), batch_time=batch_time, loss=
                losses, top1=top1, top5=top5))
    print(' * Prec@1 {top1.avg:.3f} Prec@5 {top5.avg:.3f}'.format(top1=top1,
        top5=top5))
    return top1.avg


class AverageMeter(object):
    """Computes and stores the average and current value"""

    def __init__(self):
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count


def accuracy(output, target, topk=(1,)):
    """Computes the precision@k for the specified values of k"""
    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)
        res.append(correct_k.mul_(100.0 / batch_size))
    return res


if __name__ == '__main__':
    main()

import argparse
import os
import shutil
import time
import torch
import torch.nn as nn
import torch.nn.parallel
import torch.optim
import torch.utils.data
import torchvision
import torchvision.transforms as transforms
import torchvision.datasets as datasets
import torchvision.models as models
import matplotlib.pyplot as plt
import numpy as np
from MobileNetV2 import MobileNetV2
import pdb
model_names = sorted(name for name in models.__dict__ if name.islower() and
    not name.startswith('__') and callable(models.__dict__[name]))
model_names.append('mobilenetv2')
parser = argparse.ArgumentParser(description='PyTorch ImageNet Training')
parser.add_argument('--data', metavar='DIR', help='path to dataset')
parser.add_argument('--arch', '-a', metavar='ARCH', default='mobilenetv2',
    choices=model_names, help='model architecture: ' + ' | '.join(
    model_names) + ' (default: resnet18)')
parser.add_argument('-j', '--workers', default=4, type=int, metavar='N',
    help='number of data loading workers (default: 4)')
parser.add_argument('--epochs', default=90, type=int, metavar='N', help=
    'number of total epochs to run')
parser.add_argument('--start-epoch', default=0, type=int, metavar='N', help
    ='manual epoch number (useful on restarts)')
parser.add_argument('-b', '--batch-size', default=4, type=int, metavar='N',
    help='input batch size for training (default: 20')
parser.add_argument('--lr', '--learning-rate', default=0.1, type=float,
    metavar='LR', help='initial learning rate')
parser.add_argument('--momentum', default=0.9, type=float, metavar='M',
    help='momentum')
parser.add_argument('--weight-decay', '--wd', default=0.0001, type=float,
    metavar='W', help='weight decay (default: 1e-4)')
parser.add_argument('--print-freq', '-p', default=10, type=int, metavar='N',
    help='print frequency (default: 10)')
parser.add_argument('--resume', default='', type=str, metavar='PATH', help=
    'path to latest checkpoint (default: none)')
parser.add_argument('--evaluate', default='', type=str, metavar='PATH',
    help='path to evaluate model (default: none)')
parser.add_argument('--pretrained', dest='pretrained', action='store_true',
    help='use pre-trained model')
best_prec1 = 0
MEAN_COEF = [0.48502, 0.45796, 0.4076]
DIV_COEF = [0.23068, 0.23068, 0.23068]
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize
    (mean=MEAN_COEF, std=DIV_COEF)])
train_data = datasets.CIFAR10('data', train=True, download=True, transform=
    transform)
test_data = datasets.CIFAR10('data', train=False, download=True, transform=
    transform)


def imshow(img):
    npimg = img * np.array(DIV_COEF).reshape([3, 1, 1]) + np.array(MEAN_COEF
        ).reshape([3, 1, 1])
    plt.imshow(np.transpose(npimg, (1, 2, 0)))


def main():
    global args, best_prec1
    args = parser.parse_args()
    if args.pretrained:
        print("=> using pre-trained model '{}'".format(args.arch))
        model = models.__dict__[args.arch](pretrained=True)
    else:
        print("=> creating model '{}'".format(args.arch))
        if args.arch.startswith('mobilenetv1'):
            model = MobileNetV1()
        elif args.arch.startswith('mobilenetv2'):
            model = MobileNetV2()
        else:
            model = models.__dict__[args.arch]()
        print(model)
    if args.arch.startswith('alexnet') or args.arch.startswith('vgg'):
        model.features = torch.nn.DataParallel(model.features)
        model
    else:
        model = model
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.SGD(model.parameters(), args.lr, momentum=args.
        momentum, weight_decay=args.weight_decay)
    if args.evaluate:
        if os.path.isfile(args.evaluate):
            print("=> loading model '{}'".format(args.evaluate))
            model.load_state_dict(torch.load(args.evaluate))
        else:
            print("=> no model found at '{}'".format(args.evaluate))
    val_loader = torch.utils.data.DataLoader(train_data, batch_size=args.
        batch_size, shuffle=False, num_workers=args.workers, pin_memory=True)
    if args.evaluate:
        validate(val_loader, model, criterion)
        return


def validate(val_loader, model, criterion):
    batch_time = AverageMeter()
    losses = AverageMeter()
    top1 = AverageMeter()
    top5 = AverageMeter()
    model.eval()
    end = time.time()
    for i, (inputs, target) in enumerate(val_loader):
        permute = [2, 1, 0]
        inputs = inputs[:, permute]
        target = target
        inputs = inputs
        input_var = torch.autograd.Variable(inputs, volatile=True)
        target_var = torch.autograd.Variable(target, volatile=True)
        output = model(input_var)
        loss = criterion(output.squeeze(), target_var.long())
        prec1, prec5 = accuracy(output.squeeze().data, target, topk=(1, 5))
        losses.update(loss.data[0], inputs.size(0))
        top1.update(prec1[0], inputs.size(0))
        top5.update(prec5[0], inputs.size(0))
        batch_time.update(time.time() - end)
        end = time.time()
        if i % args.print_freq == 0:
            print(
                'Test: [{0}/{1}]\tTime {batch_time.val:.3f} ({batch_time.avg:.3f})\tLoss {loss.val:.4f} ({loss.avg:.4f})\tPrec@1 {top1.val:.3f} ({top1.avg:.3f})\tPrec@5 {top5.val:.3f} ({top5.avg:.3f})'
                .format(i, len(val_loader), batch_time=batch_time, loss=
                losses, top1=top1, top5=top5))
    print(' * Prec@1 {top1.avg:.3f} Prec@5 {top5.avg:.3f}'.format(top1=top1,
        top5=top5))
    return top1.avg


class AverageMeter(object):
    """Computes and stores the average and current value"""

    def __init__(self):
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count


def accuracy(output, target, topk=(1,)):
    """Computes the precision@k for the specified values of k"""
    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)
        res.append(correct_k.mul_(100.0 / batch_size))
    return res


if __name__ == '__main__':
    main()

import argparse
import os
#os.environ['CUDA_VISIBLE_DEVICES']='2'
import shutil
import time

import torch
import torch.nn as nn
import torch.nn.parallel
#import torch.backends.cudnn as cudnn
import torch.optim
import torch.utils.data
import torchvision
import torchvision.transforms as transforms
import torchvision.datasets as datasets
import torchvision.models as models
import matplotlib.pyplot as plt
import numpy as np
from MobileNetV2 import MobileNetV2
import pdb


model_names = sorted(name for name in models.__dict__
    if name.islower() and not name.startswith("__")
    and callable(models.__dict__[name]))
model_names.append('mobilenetv2')


parser = argparse.ArgumentParser(description='PyTorch ImageNet Training')
parser.add_argument('--data', metavar='DIR',
                    help='path to dataset')
parser.add_argument('--arch', '-a', metavar='ARCH', default='mobilenetv2',
                    choices=model_names,
                    help='model architecture: ' +
                        ' | '.join(model_names) +
                        ' (default: resnet18)')
parser.add_argument('-j', '--workers', default=4, type=int, metavar='N',
                    help='number of data loading workers (default: 4)')
parser.add_argument('--epochs', default=90, type=int, metavar='N',
                    help='number of total epochs to run')
parser.add_argument('--start-epoch', default=0, type=int, metavar='N',
                    help='manual epoch number (useful on restarts)')
parser.add_argument('-b', '--batch-size', default=4, type=int,
                    metavar='N',help='input batch size for training (default: 20')
parser.add_argument('--lr', '--learning-rate', default=0.1, type=float,
                    metavar='LR', help='initial learning rate')
parser.add_argument('--momentum', default=0.9, type=float, metavar='M',
                    help='momentum')
parser.add_argument('--weight-decay', '--wd', default=1e-4, type=float,
                    metavar='W', help='weight decay (default: 1e-4)')
parser.add_argument('--print-freq', '-p', default=10, type=int,
                    metavar='N', help='print frequency (default: 10)')
parser.add_argument('--resume', default='', type=str, metavar='PATH',
                    help='path to latest checkpoint (default: none)')
parser.add_argument('--evaluate', default='', type=str, metavar='PATH',
                    help='path to evaluate model (default: none)')
parser.add_argument('--pretrained', dest='pretrained', action='store_true',
                    help='use pre-trained model')

best_prec1 = 0


MEAN_COEF = [0.48502, 0.45796, 0.40760]
DIV_COEF = [0.23068, 0.23068, 0.23068]

transform=transforms.Compose([transforms.ToTensor(),transforms.Normalize(mean=MEAN_COEF,std=DIV_COEF)])
train_data=datasets.CIFAR10('data',train=True,download=True,transform=transform)
test_data=datasets.CIFAR10('data',train=False,download=True,transform=transform)

# functions to show an image

def imshow(img):
    npimg = img * np.array(DIV_COEF).reshape([3,1,1]) \
          + np.array(MEAN_COEF).reshape([3,1,1])            # Un-normalize
    plt.imshow(np.transpose(npimg, (1, 2, 0)))


def main():
    global args, best_prec1
    args = parser.parse_args()

    # create model
    if args.pretrained:
        print("=> using pre-trained model '{}'".format(args.arch))
        model = models.__dict__[args.arch](pretrained=True)
    else:
        print("=> creating model '{}'".format(args.arch))
        if args.arch.startswith('mobilenetv1'):
            model = MobileNetV1()
        elif args.arch.startswith('mobilenetv2'):
            model = MobileNetV2()
        else:
            model = models.__dict__[args.arch]()
        print(model)

    if args.arch.startswith('alexnet') or args.arch.startswith('vgg'):
        model.features = torch.nn.DataParallel(model.features)
        model
    else:
         model=model

    # define loss function (criterion) and optimizer
    criterion = nn.CrossEntropyLoss()

    optimizer = torch.optim.SGD(model.parameters(), args.lr,
                                momentum=args.momentum,
                                weight_decay=args.weight_decay)


    if args.evaluate:
        if os.path.isfile(args.evaluate):
            print("=> loading model '{}'".format(args.evaluate))
            model.load_state_dict(torch.load(args.evaluate))
        else:
            print("=> no model found at '{}'".format(args.evaluate))


    #cudnn.benchmark = True

    val_loader = torch.utils.data.DataLoader(train_data,
        batch_size=args.batch_size,
        shuffle=False,
        num_workers=args.workers, pin_memory=True)
    

    if args.evaluate:
        validate(val_loader, model, criterion)
        return


def validate(val_loader, model, criterion):
    batch_time = AverageMeter()
    losses = AverageMeter()
    top1 = AverageMeter()
    top5 = AverageMeter()

    # switch to evaluate mode
    model.eval()

    end = time.time()
    for i, (inputs, target) in enumerate(val_loader):
        #pdb.set_trace()
        permute = [2, 1, 0]
        inputs = inputs[:,permute]

        #imshow(torchvision.utils.make_grid(inputs))
        #plt.show()


        target = target
        inputs = inputs
        input_var = torch.autograd.Variable(inputs, volatile=True)
        target_var = torch.autograd.Variable(target, volatile=True)
        
        # compute output
        output = model(input_var)

        
        #pdb.set_trace()
        loss = criterion(output.squeeze(), target_var.long())

        # measure accuracy and record loss
        prec1, prec5 = accuracy(output.squeeze().data, target, topk=(1, 5))
        losses.update(loss.data[0], inputs.size(0))
        top1.update(prec1[0], inputs.size(0))
        top5.update(prec5[0], inputs.size(0))

        # measure elapsed time
        batch_time.update(time.time() - end)
        end = time.time()

        if i % args.print_freq == 0:
            print('Test: [{0}/{1}]\t'
                  'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
                  'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
                  'Prec@1 {top1.val:.3f} ({top1.avg:.3f})\t'
                  'Prec@5 {top5.val:.3f} ({top5.avg:.3f})'.format(
                   i, len(val_loader), batch_time=batch_time, loss=losses,
                   top1=top1, top5=top5))

    print(' * Prec@1 {top1.avg:.3f} Prec@5 {top5.avg:.3f}'
          .format(top1=top1, top5=top5))



    
    return top1.avg


class AverageMeter(object):
    """Computes and stores the average and current value"""
    def __init__(self):
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count


def accuracy(output, target, topk=(1,)):
    """Computes the precision@k for the specified values of k"""
    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)
        res.append(correct_k.mul_(100.0 / batch_size))
    return res


if __name__ == '__main__':
    main()