help="nb_grads parameter for DeepFool adversarial trainer")
args = parser.parse_args()
## Step 1: load model and dataset
# load the base classifier
checkpoint = torch.load(args.base_classifier)
base_classifier = get_architecture(checkpoint["arch"])
base_classifier.load_state_dict(checkpoint['model_state_dict'])
# create the smooothed classifier g
smoothed_classifier = Smooth(base_classifier, get_num_classes(), args.sigma)
# # iterate through the dataset
dataset = datasets.CIFAR10("./dataset_cache",
train=False,
download=True,
transform=transforms.ToTensor())
y_test = np.asarray(dataset.targets) # test labels
min_pixel_value = 0.0 # min value
max_pixel_value = 1.0 # max value
# Step 2: create an interface for classifier, load trained model, to be used by attack model trainer
# Define the loss function and the optimizer for attack model trainer
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(base_classifier.parameters(),
lr=0.1,
momentum=0.9,
weight_decay=1e-4)
def main():
# Init logger6
if not os.path.isdir(args.save_path):
os.makedirs(args.save_path)
log = open(
os.path.join(args.save_path,
'log_seed_{}.txt'.format(args.manualSeed)), 'w')
print_log('save path : {}'.format(args.save_path), log)
state = {k: v for k, v in args._get_kwargs()}
print_log(state, log)
print_log("Random Seed: {}".format(args.manualSeed), log)
print_log("python version : {}".format(sys.version.replace('\n', ' ')),
log)
print_log("torch version : {}".format(torch.__version__), log)
print_log("cudnn version : {}".format(torch.backends.cudnn.version()),
log)
print_log("Weight Decay: {}".format(args.decay), log)
# Init the tensorboard path and writer
tb_path = os.path.join(args.save_path, 'tb_log')
logger = Logger(tb_path)
writer = SummaryWriter(tb_path)
# Init dataset
if not os.path.isdir(args.data_path):
os.makedirs(args.data_path)
if args.dataset == 'cifar10':
mean = [x / 255 for x in [125.3, 123.0, 113.9]]
std = [x / 255 for x in [63.0, 62.1, 66.7]]
elif args.dataset == 'cifar100':
mean = [x / 255 for x in [129.3, 124.1, 112.4]]
std = [x / 255 for x in [68.2, 65.4, 70.4]]
elif args.dataset == 'svhn':
mean = [0.5, 0.5, 0.5]
std = [0.5, 0.5, 0.5]
elif args.dataset == 'mnist':
mean = [0.5, 0.5, 0.5]
std = [0.5, 0.5, 0.5]
elif args.dataset == 'imagenet':
mean = [0.485, 0.456, 0.406]
std = [0.229, 0.224, 0.225]
elif args.dataset == 'yawnDD':
print(f'YawnDD dataset!')
elif args.dataset == 'eyeclosure':
print(f'eyeclosure dataset!')
else:
assert False, "Unknow dataset : {}".format(args.dataset)
if args.dataset == 'imagenet':
train_transform = transforms.Compose([
transforms.RandomResizedCrop(224, scale=(0.2, 1.0)),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize(mean, std)
])
test_transform = transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize(mean, std)
]) # here is actually the validation dataset
elif not args.dataset in ['yawnDD', 'eyeclosure']:
train_transform = transforms.Compose([
transforms.RandomHorizontalFlip(),
transforms.RandomCrop(32, padding=4),
transforms.ToTensor(),
transforms.Normalize(mean, std)
])
test_transform = transforms.Compose([
transforms.RandomCrop(32, padding=4),
transforms.ToTensor(),
transforms.Normalize(mean, std)
])
if args.dataset == 'mnist':
train_data = dset.MNIST(args.data_path,
train=True,
transform=train_transform,
download=True)
test_data = dset.MNIST(args.data_path,
train=False,
transform=test_transform,
download=True)
num_classes = 10
elif args.dataset == 'cifar10':
train_data = dset.CIFAR10(args.data_path,
train=True,
transform=train_transform,
download=True)
test_data = dset.CIFAR10(args.data_path,
train=False,
transform=test_transform,
download=True)
num_classes = 10
elif args.dataset == 'cifar100':
train_data = dset.CIFAR100(args.data_path,
train=True,
transform=train_transform,
download=True)
test_data = dset.CIFAR100(args.data_path,
train=False,
transform=test_transform,
download=True)
num_classes = 100
elif args.dataset == 'svhn':
train_data = dset.SVHN(args.data_path,
split='train',
transform=train_transform,
download=True)
test_data = dset.SVHN(args.data_path,
split='test',
transform=test_transform,
download=True)
num_classes = 10
elif args.dataset == 'stl10':
train_data = dset.STL10(args.data_path,
split='train',
transform=train_transform,
download=True)
test_data = dset.STL10(args.data_path,
split='test',
transform=test_transform,
download=True)
num_classes = 10
elif args.dataset == 'imagenet':
train_dir = os.path.join(args.data_path, 'train')
test_dir = os.path.join(args.data_path, 'val')
train_data = dset.ImageFolder(train_dir, transform=train_transform)
test_data = dset.ImageFolder(test_dir, transform=test_transform)
num_classes = 1000
elif args.dataset == 'yawnDD':
dataset = torch.load('./yawning_dataset/yawnDD_image.pt') / 255
target = torch.load('./yawning_dataset/yawnDD_label.pt').long() / 255
dataset = dataset.view(dataset.size(0), dataset.size(3),
dataset.size(2), dataset.size(2))
train_dataset = dataset[:int(0.8 * dataset.size(0))]
train_target = target[:int(0.8 * dataset.size(0))]
test_dataset = dataset[-int(0.2 * dataset.size(0)):]
test_target = target[-int(0.2 * dataset.size(0)):]
train_data = torch.utils.data.TensorDataset(train_dataset,
train_target)
test_data = torch.utils.data.TensorDataset(test_dataset, test_target)
num_classes = 2
elif args.dataset == 'eyeclosure':
train_dataset = torch.load('./eyeclosure/eyeclosure_train_data.pt')
train_label = torch.load(
'./eyeclosure/eyeclosure_train_label.pt').long()
test_dataset = torch.load('./eyeclosure/eyeclosure_test_data.pt')
test_label = torch.load('./eyeclosure/eyeclosure_test_label.pt').long()
train_dataset = train_dataset.view(train_dataset.size(0),
train_dataset.size(3),
train_dataset.size(2),
train_dataset.size(2))
test_dataset = test_dataset.view(test_dataset.size(0),
test_dataset.size(3),
test_dataset.size(2),
test_dataset.size(2))
train_data = torch.utils.data.TensorDataset(train_dataset, train_label)
test_data = torch.utils.data.TensorDataset(test_dataset, test_label)
num_classes = 2
else:
assert False, 'Do not support dataset : {}'.format(args.dataset)
train_loader = torch.utils.data.DataLoader(train_data,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=True)
test_loader = torch.utils.data.DataLoader(test_data,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
print_log("=> creating model '{}'".format(args.arch), log)
# Init model, criterion, and optimizer
net = models.__dict__[args.arch](num_classes)
print_log("=> network :\n {}".format(net), log)
if args.use_cuda:
if args.ngpu > 1:
# net = torch.nn.DataParallel(net, device_ids=[1,2])
net = torch.nn.DataParallel(net)
# define loss function (criterion) and optimizer
criterion = torch.nn.CrossEntropyLoss()
if args.optimizer == "SGD":
print("using SGD as optimizer")
optimizer = torch.optim.SGD(filter(lambda param: param.requires_grad,
net.parameters()),
lr=state['learning_rate'],
momentum=state['momentum'],
weight_decay=state['decay'],
nesterov=True)
elif args.optimizer == "Adam":
print("using Adam as optimizer")
optimizer = torch.optim.Adam(filter(lambda param: param.requires_grad,
net.parameters()),
lr=state['learning_rate'],
weight_decay=state['decay'])
elif args.optimizer == "YF":
print("using YellowFin as optimizer")
optimizer = YFOptimizer(filter(lambda param: param.requires_grad,
net.parameters()),
lr=state['learning_rate'],
mu=state['momentum'],
weight_decay=state['decay'])
elif args.optimizer == "RMSprop":
print("using RMSprop as optimizer")
optimizer = torch.optim.RMSprop(filter(
lambda param: param.requires_grad, net.parameters()),
lr=state['learning_rate'],
alpha=0.99,
eps=1e-08,
weight_decay=0,
momentum=0)
if args.use_cuda:
net.cuda()
criterion.cuda()
recorder = RecorderMeter(args.epochs) # count number of epoches
for name, value in net.named_parameters():
print(name)
# optionally resume from a checkpoint
if args.resume:
new_state_dict = OrderedDict()
if os.path.isfile(args.resume):
print_log("=> loading checkpoint '{}'".format(args.resume), log)
checkpoint = torch.load(args.resume)
if not (args.fine_tune):
args.start_epoch = checkpoint['epoch']
recorder = checkpoint['recorder']
optimizer.load_state_dict(checkpoint['optimizer'])
state_tmp = net.state_dict()
for k, v in checkpoint['state_dict'].items():
print(k)
name = k
# name = k[7:]
print(name)
new_state_dict[name] = v
if 'state_dict' in checkpoint.keys():
#state_tmp.update(new_state_dict['state_dict'])
state_tmp.update(new_state_dict)
else:
print('loading from pth file not tar file')
state_tmp.update(new_state_dict)
#state_tmp.update(checkpoint)
net.load_state_dict(state_tmp)
#net.load_state_dict(torch.load('save/mobilenetv2_1.pth'))
# net.load_state_dict(checkpoint['state_dict'])
print_log(
"=> loaded checkpoint '{}' (epoch {})".format(
args.resume, args.start_epoch), log)
else:
print_log("=> no checkpoint found at '{}'".format(args.resume),
log)
else:
print_log(
"=> do not use any checkpoint for {} model".format(args.arch), log)
if args.evaluate:
validate(test_loader, net, criterion, log)
return
# Main loop
start_time = time.time()
epoch_time = AverageMeter()
for epoch in range(args.start_epoch, args.epochs):
current_learning_rate, current_momentum = adjust_learning_rate(
optimizer, epoch, args.gammas, args.schedule)
# Display simulation time
need_hour, need_mins, need_secs = convert_secs2time(
epoch_time.avg * (args.epochs - epoch))
need_time = '[Need: {:02d}:{:02d}:{:02d}]'.format(
need_hour, need_mins, need_secs)
print_log(
'\n==>>{:s} [Epoch={:03d}/{:03d}] {:s} [LR={:6.5f}][M={:1.2f}]'.
format(time_string(), epoch, args.epochs, need_time,
current_learning_rate, current_momentum) +
' [Best : Accuracy={:.2f}, Error={:.2f}]'.format(
recorder.max_accuracy(False),
100 - recorder.max_accuracy(False)), log)
# train for one epoch
train_acc, train_los = train(train_loader, net, criterion, optimizer,
epoch, log)
# evaluate on validation set
val_acc, val_los = validate(test_loader, net, criterion, log)
is_best = val_acc > recorder.max_accuracy(istrain=False)
recorder.update(epoch, train_los, train_acc, val_los, val_acc)
if args.model_only:
checkpoint_state = {'state_dict': net.state_dict()}
else:
checkpoint_state = {
'epoch': epoch + 1,
'arch': args.arch,
'state_dict': net.state_dict(),
'recorder': recorder,
'optimizer': optimizer.state_dict(),
}
save_checkpoint(checkpoint_state, is_best, args.save_path,
f'checkpoint.pth.tar', log)
# measure elapsed time
epoch_time.update(time.time() - start_time)
start_time = time.time()
recorder.plot_curve(os.path.join(args.save_path, 'curve.png'))
log.close()
import torch.optim as optim
import torch.utils.data
import torchvision.datasets as dset
import torchvision.transforms as transforms
import torchvision.utils as vutils
from torch.autograd import Variable
# Setting some hyperparameters
batchSize = 64 # We set the size of the batch.
imageSize = 64 # We set the size of the generated images (64x64).
# Creating the transformations
transform = transforms.Compose([transforms.Scale(imageSize), transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),]) # We create a list of transformations (scaling, tensor conversion, normalization) to apply to the input images.
# Loading the dataset
dataset = dset.CIFAR10(root = './data', download = True, transform = transform) # We download the training set in the ./data folder and we apply the previous transformations on each image.
dataloader = torch.utils.data.DataLoader(dataset, batch_size = batchSize, shuffle = True, num_workers = 0) # We use dataLoader to get the images of the training set batch by batch.
# Defining the weights_init function that takes as input a neural network m and that will initialize all its weights.
def weights_init(m):
classname = m.__class__.__name__
if classname.find('Conv') != -1:
m.weight.data.normal_(0.0, 0.02)
elif classname.find('BatchNorm') != -1:
m.weight.data.normal_(1.0, 0.02)
m.bias.data.fill_(0)
# Defining the generator
class G(nn.Module): # We introduce a class to define the generator.
import torch
from torch import nn
from torchvision import datasets, transforms, models
from torch import nn
# net = models.resnet18()
# print(net)
# # m = nn.ZeroPad2d(16)
# # x = torch.randn(1, 3, 32, 32)
# # print(m(x).shape)
train_data = datasets.CIFAR10("./data",
True,
transform=transforms.ToTensor(),
download=True)
print(train_data[0][0].shape)
from __future__ import print_function
from argparse import ArgumentParser
from torch.utils.data import DataLoader
import torchvision.datasets as dset
from torchvision import datasets, transforms
import torch
tsfm=transforms.Compose([
transforms.ToTensor()
])
dataset = dset.CIFAR10(root='./data',train=False, download=True,transform=tsfm)
data_loader = torch.utils.data.DataLoader(dataset,batch_size=2000, num_workers=16)
print(len(data_loader))
for i,data in enumerate(data_loader):
if i==0:
imgs,labels = torch.round(data[0]*255),data[1]
else:
imgs,labels = torch.cat([imgs,torch.round(data[0]*255)],dim=0),torch.cat([labels,data[1]],dim=0)
print(imgs.size(),imgs.max(),imgs.min(),labels.size())
torch.save([imgs,labels],'test.pt')
# !/usr/bin/env python
# -*- coding: utf-8 -*-
train_dataset = datasets.MNIST(args.data_dir,
train=True,
download=False,
transform=train_transform)
test_dataset = datasets.MNIST(args.data_dir,
train=False,
download=False,
transform=test_transform)
elif args.model == 'AlexNet':
model = AlexNetForCIFAR()
train_transform, test_transform = get_data_transform('cifar')
train_dataset = datasets.CIFAR10(args.data_dir,
train=True,
download=False,
transform=train_transform)
test_dataset = datasets.CIFAR10(args.data_dir,
train=False,
download=False,
transform=test_transform)
elif args.model == 'LROnMnist':
model = ResNetOnCifar10.LROnMnist()
train_transform, test_transform = get_data_transform('mnist')
train_dataset = datasets.MNIST(args.data_dir,
train=True,
download=False,
transform=train_transform)
test_dataset = datasets.MNIST(args.data_dir,
train=False,
def get_datasets(args, cutout):
name = args.dataset
root = args.data_path
if name == 'cifar10':
mean = [x / 255 for x in [125.3, 123.0, 113.9]]
std = [x / 255 for x in [63.0, 62.1, 66.7]]
elif name == 'cifar100':
mean = [x / 255 for x in [129.3, 124.1, 112.4]]
std = [x / 255 for x in [68.2, 65.4, 70.4]]
elif name.startswith('imagenet-1k'):
mean, std = [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]
elif name.startswith('ImageNet16'):
mean = [x / 255 for x in [122.68, 116.66, 104.01]]
std = [x / 255 for x in [63.22, 61.26, 65.09]]
elif name in [
'MiniImageNet', 'MetaMiniImageNet', 'TieredImageNet',
'MetaTieredImageNet'
]:
pass
else:
raise TypeError("Unknow dataset : {:}".format(name))
# Data Argumentation
if name == 'cifar10' or name == 'cifar100':
lists = [
transforms.RandomHorizontalFlip(),
transforms.RandomCrop(32, padding=4),
transforms.ToTensor(),
transforms.Normalize(mean, std)
]
if cutout > 0: lists += [CUTOUT(cutout)]
train_transform = transforms.Compose(lists)
test_transform = transforms.Compose(
[transforms.ToTensor(),
transforms.Normalize(mean, std)])
xshape = (1, 3, 32, 32)
elif name.startswith('ImageNet16'):
lists = [
transforms.RandomHorizontalFlip(),
transforms.RandomCrop(16, padding=2),
transforms.ToTensor(),
transforms.Normalize(mean, std)
]
if cutout > 0: lists += [CUTOUT(cutout)]
train_transform = transforms.Compose(lists)
test_transform = transforms.Compose(
[transforms.ToTensor(),
transforms.Normalize(mean, std)])
xshape = (1, 3, 16, 16)
elif name == 'tiered':
lists = [
transforms.RandomHorizontalFlip(),
transforms.RandomCrop(80, padding=4),
transforms.ToTensor(),
transforms.Normalize(mean, std)
]
if cutout > 0: lists += [CUTOUT(cutout)]
train_transform = transforms.Compose(lists)
test_transform = transforms.Compose([
transforms.CenterCrop(80),
transforms.ToTensor(),
transforms.Normalize(mean, std)
])
xshape = (1, 3, 32, 32)
elif name.startswith('imagenet-1k'):
lists = [
transforms.Resize((32, 32), interpolation=2),
transforms.RandomHorizontalFlip(),
transforms.RandomCrop(32, padding=4),
transforms.ToTensor(),
transforms.Normalize(mean, std)
]
if cutout > 0: lists += [CUTOUT(cutout)]
train_transform = transforms.Compose(lists)
test_transform = transforms.Compose(
[transforms.ToTensor(),
transforms.Normalize(mean, std)])
xshape = (1, 3, 32, 32)
elif name in [
'MiniImageNet', 'MetaMiniImageNet', 'TieredImageNet',
'MetaTieredImageNet'
]:
train_transform, test_transform = transforms_options[args.transform]
if args.transform in ['A', 'B']:
xshape = (1, 3, 84, 84)
elif args.transform == 'C':
xshape = (1, 3, 80, 80)
elif args.transform == 'D':
xshape = (1, 3, 32, 32)
else:
raise TypeError("Unknow dataset : {:}".format(name))
if name == 'cifar10':
train_data = dset.CIFAR10(root,
train=True,
transform=train_transform,
download=True)
test_data = dset.CIFAR10(root,
train=False,
transform=test_transform,
download=True)
assert len(train_data) == 50000 and len(test_data) == 10000
elif name == 'cifar100':
train_data = dset.CIFAR100(root,
train=True,
transform=train_transform,
download=True)
test_data = dset.CIFAR100(root,
train=False,
transform=test_transform,
download=True)
assert len(train_data) == 50000 and len(test_data) == 10000
elif name.startswith('imagenet-1k'):
train_data = dset.ImageFolder(osp.join(root, 'train'), train_transform)
test_data = dset.ImageFolder(osp.join(root, 'val'), test_transform)
elif name == 'ImageNet16':
train_data = ImageNet16(root, True, train_transform)
test_data = ImageNet16(root, False, test_transform)
assert len(train_data) == 1281167 and len(test_data) == 50000
elif name == 'ImageNet16-120':
train_data = ImageNet16(root, True, train_transform, 120)
test_data = ImageNet16(root, False, test_transform, 120)
assert len(train_data) == 151700 and len(test_data) == 6000
elif name == 'ImageNet16-150':
train_data = ImageNet16(root, True, train_transform, 150)
test_data = ImageNet16(root, False, test_transform, 150)
assert len(train_data) == 190272 and len(test_data) == 7500
elif name == 'ImageNet16-200':
train_data = ImageNet16(root, True, train_transform, 200)
test_data = ImageNet16(root, False, test_transform, 200)
assert len(train_data) == 254775 and len(test_data) == 10000
elif name == 'MiniImageNet':
train_data = ImageNet(args=args,
partition='train',
transform=train_transform)
test_data = ImageNet(args=args,
partition='val',
transform=test_transform)
elif name == 'MetaMiniImageNet':
train_data = MetaImageNet(args=args,
partition='test',
train_transform=train_transform,
test_transform=test_transform)
test_data = MetaImageNet(args=args,
partition='val',
train_transform=train_transform,
test_transform=test_transform)
elif name == 'TieredImageNet':
train_data = TieredImageNet(args=args,
partition='train',
transform=train_transform)
test_data = TieredImageNet(args=args,
partition='train_phase_val',
transform=test_transform)
elif name == 'MetaTieredImageNet':
train_data = MetaTieredImageNet(args=args,
partition='test',
train_transform=train_transform,
test_transform=test_transform)
test_data = MetaTieredImageNet(args=args,
partition='val',
train_transform=train_transform,
test_transform=test_transform)
else:
raise TypeError("Unknow dataset : {:}".format(name))
class_num = Dataset2Class[name]
return train_data, test_data, xshape, class_num
def main():
#SAVE filename
path_current = os.path.dirname(os.path.realpath(__file__))
path_subdir = 'dataset'
data_filename = 'VGGdataset.txt'
path_file = os.path.join(path_current,path_subdir,data_filename)
f=open(path_file,'w')
global args, best_prec1
args = parser.parse_args()
use_cuda = not args.no_cuda and torch.cuda.is_available()
# 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))
# model = models.__dict__[args.arch]()
#
# if not args.distributed:
# if args.arch.startswith('alexnet') or args.arch.startswith('vgg'):
# model.features = torch.nn.DataParallel(model.features)
# model.cuda()
# else:
# model = torch.nn.DataParallel(model).cuda()
# else:
# model.cuda()
# model = torch.nn.parallel.DistributedDataParallel(model)
#
device = torch.device("cuda" if use_cuda else "cpu")
# model = vggNet().cuda().to("cuda")
model = vggNet(make_conv_layers(vgg16_config),make_fc_layers()).cuda()
def init_vgg(net):
import math
for m in net.modules():
if isinstance(m, nn.Conv2d):
c_out, _, kh, kw = m.weight.size()
n = kh * kw * c_out
m.weight.data.normal_(0, math.sqrt(2 / n))
if m.bias is not None:
m.bias.data.zero_()
elif isinstance(m, nn.BatchNorm2d):
m.weight.data.fill_(1)
m.bias.data.zero_()
elif isinstance(m, nn.Linear):
m.weight.data.normal_(0, 0.01)
m.bias.data.zero_()
init_vgg(model)
# define loss function (criterion) and optimizer
criterion = nn.CrossEntropyLoss().cuda()
optimizer = torch.optim.SGD(model.parameters(), args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
# optionally resume from a checkpoint
if args.resume:
if os.path.isfile(args.resume):
print("=> loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch']
best_prec1 = checkpoint['best_prec1']
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
print("=> loaded checkpoint '{}' (epoch {})"
.format(args.resume, checkpoint['epoch']))
else:
print("=> no checkpoint found at '{}'".format(args.resume))
cudnn.benchmark = True
# Data loading code
# traindir = os.path.join(args.data, 'train')
# valdir = os.path.join(args.data, 'val')
# normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
# std=[0.229, 0.224, 0.225])
#
# train_dataset = datasets.ImageFolder(
# traindir,
# transforms.Compose([
# transforms.RandomResizedCrop(224),
# transforms.RandomHorizontalFlip(),
# transforms.ToTensor(),
# normalize,
# ]))
#
#
#
# if args.distributed:
# train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
# else:
# train_sampler = None
train_loader = torch.utils.data.DataLoader(
datasets.CIFAR10(
'../data',
train=True,
download=True,
transform=transforms.Compose([
transforms.RandomHorizontalFlip(),
transforms.RandomCrop(32,4),
# transforms.ColorJitter(),
transforms.ToTensor(),
transforms.Normalize(
mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]
),
])
),
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=True
)
# train_loader = torch.utils.data.DataLoader(
# train_dataset, batch_size=args.batch_size, shuffle=(train_sampler is None),
# num_workers=args.workers, pin_memory=True, sampler=train_sampler)
val_loader = torch.utils.data.DataLoader(
datasets.CIFAR10(
'../data',
train=False,
download=True,
transform = transforms.Compose([
# transforms.RandomHorizontalFlip(),
# transforms.ColorJitter(),
transforms.ToTensor(),
transforms.Normalize(
mean = [0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]
),
])
),
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True
)
# val_loader = torch.utils.data.DataLoader(
# datasets.ImageFolder(valdir, transforms.Compose([
# transforms.Resize(256),
# transforms.CenterCrop(224),
# transforms.ToTensor(),
# normalize,
# ])),
# batch_size=args.batch_size, shuffle=False,
# num_workers=args.workers, pin_memory=True)
# if args.evaluate:
# validate(val_loader, model, criterion)
# return
train_losses =np.zeros((args.epochs))
train_prec1s =np.zeros((args.epochs))
eval_losses =np.zeros((args.epochs))
eval_prec1s =np.zeros((args.epochs))
x_epoch = np.zeros((args.epochs))
# x_epoch = np.linspace(args.start_epoch,args.epochs -1, args.epochs,endpoint=True)
for epoch in range(args.start_epoch, args.epochs):
adjust_learning_rate(optimizer, epoch)
# train for one epoch
train_loss, train_prec1 = train(train_loader, model, criterion, optimizer, epoch,f)
# evaluate on validation set
eval_loss, eval_prec1 = validate(val_loader, model, criterion,f)
train_losses[epoch] = train_loss
train_prec1s[epoch] = train_prec1
eval_losses[epoch] = eval_loss
eval_prec1s[epoch] = eval_prec1
x_epoch[epoch] = epoch
# train_losses =np.append( train_losses + train_loss
# train_prec1s = train_prec1s + train_prec1
# eval_losses = eval_losses + eval_loss
# eval_prec1s = eval_prec1s + eval_prec1
##
# train_loss.append(train_losses)
# train_prec1.append(train_prec1s)
# eval_loss.append(eval_losses)
# eval_prec1.append(eval_prec1s)
# remember best prec@1 and save checkpoint
is_best = eval_prec1 > best_prec1
best_prec1 = max(eval_prec1, best_prec1)
save_checkpoint({
'epoch': epoch + 1,
'arch': 'VGG16',
'state_dict': model.state_dict(),
'best_prec1': best_prec1,
'optimizer' : optimizer.state_dict(),
}, is_best)
matplotlib.use('Agg')
plt.clf()
plt.close()
fig_loss = plt.figure()
ax_loss = fig_loss.add_subplot(1,1,1)
ax_loss.plot(x_epoch,train_losses,label='Train Loss')
ax_loss.plot(x_epoch,eval_losses,label='Test Loss')
ax_loss.legend(loc=1)
ax_loss.set_xlabel('epoch')
ax_loss.set_ylabel('loss')
ax_loss.set_title('Loss of Train and Test')
plot_loss_filename = 'VGGloss.png'
path_loss_file = os.path.join(path_current,path_subdir,plot_loss_filename)
fig_loss.savefig(path_loss_file)
plt.clf()
plt.close()
fig_prec = plt.figure()
ax_prec = fig_prec.add_subplot(1,1,1)
ax_prec.plot(x_epoch,train_prec1s,label='Train Best1')
ax_prec.plot(x_epoch,eval_prec1s,label='Test Best1')
ax_prec.legend(loc=1)
ax_prec.set_xlabel('epoch')
ax_prec.set_ylabel('Best1 Precision')
ax_prec.set_title('Best1 of Train and Test')
plot_prec_filename = 'VGGprec.png'
path_prec_file = os.path.join(path_current,path_subdir,plot_prec_filename)
fig_prec.savefig(path_prec_file)
f.close()
train=True,
download=True,
transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307, ),
(0.3081, ))
]))
img_size = dataset_train[0][0].shape
elif args.dataset == 'cifar':
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
dataset_train = datasets.CIFAR10('./data/cifar',
train=True,
transform=transform,
target_transform=None,
download=True)
img_size = dataset_train[0][0].shape
else:
exit('Error: unrecognized dataset')
# testing
if args.dataset == 'mnist':
dataset_test = datasets.MNIST('./data/mnist/',
train=False,
download=True,
transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307, ),
(0.3081, ))
def main():
# Scale and initialize the parameters
best_prec1 = 0
configs.TRAIN.epochs = int(
math.ceil(configs.TRAIN.epochs / configs.ADV.n_repeats))
configs.ADV.fgsm_step /= configs.DATA.max_color_value
configs.ADV.clip_eps /= configs.DATA.max_color_value
# Create output folder
if not os.path.isdir(os.path.join('trained_models', configs.output_name)):
os.makedirs(os.path.join('trained_models', configs.output_name))
# Log the config details
logger.info(pad_str(' ARGUMENTS '))
for k, v in configs.items():
print('{}: {}'.format(k, v))
logger.info(pad_str(''))
# Create the model
if configs.pretrained:
print("=> using pre-trained model '{}'".format(configs.TRAIN.arch))
model = models.__dict__[configs.TRAIN.arch](pretrained=True)
else:
print("=> creating model '{}'".format(configs.TRAIN.arch))
model = models.__dict__[configs.TRAIN.arch]()
# Wrap the model into DataParallel
model = torch.nn.DataParallel(model).cuda()
# Criterion:
criterion = nn.CrossEntropyLoss().cuda()
# Optimizer:
optimizer = torch.optim.SGD(model.parameters(),
configs.TRAIN.lr,
momentum=configs.TRAIN.momentum,
weight_decay=configs.TRAIN.weight_decay)
# Resume if a valid checkpoint path is provided
if configs.resume:
if os.path.isfile(configs.resume):
print("=> loading checkpoint '{}'".format(configs.resume))
checkpoint = torch.load(configs.resume)
configs.TRAIN.start_epoch = checkpoint['epoch']
best_prec1 = checkpoint['best_prec1']
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
print("=> loaded checkpoint '{}' (epoch {})".format(
configs.resume, checkpoint['epoch']))
else:
print("=> no checkpoint found at '{}'".format(configs.resume))
# Initiate data loaders
traindir = os.path.join(configs.data, 'train')
valdir = os.path.join(configs.data, 'val')
# train_dataset = datasets.ImageFolder(
# traindir,
# transforms.Compose([
# transforms.RandomResizedCrop(configs.DATA.crop_size),
# transforms.RandomHorizontalFlip(),
# transforms.ToTensor(),
# ]))
training_data = datasets.CIFAR10(root="data",
train=True,
download=True,
transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5),
(1.0, 1.0, 1.0))
]))
train_loader = torch.utils.data.DataLoader(
training_data,
batch_size=configs.DATA.batch_size,
shuffle=True,
num_workers=configs.DATA.workers,
pin_memory=True,
sampler=None)
test_data = datasets.CIFAR10(root="data",
train=False,
download=True,
transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5),
(1.0, 1.0, 1.0))
]))
normalize = transforms.Normalize(mean=configs.TRAIN.mean,
std=configs.TRAIN.std)
val_loader = torch.utils.data.DataLoader(
test_data,
batch_size=configs.DATA.batch_size,
shuffle=False,
num_workers=configs.DATA.workers,
pin_memory=True)
# If in evaluate mode: perform validation on PGD attacks as well as clean samples
if configs.evaluate:
logger.info(pad_str(' Performing PGD Attacks '))
for pgd_param in configs.ADV.pgd_attack:
validate_pgd(val_loader, model, criterion, pgd_param[0],
pgd_param[1], configs, logger)
validate(val_loader, model, criterion, configs, logger)
return
for epoch in range(configs.TRAIN.start_epoch, configs.TRAIN.epochs):
adjust_learning_rate(configs.TRAIN.lr, optimizer, epoch,
configs.ADV.n_repeats)
# train for one epoch
train(train_loader, model, criterion, optimizer, epoch)
# evaluate on validation set
prec1 = validate(val_loader, model, criterion, configs, logger)
# remember best prec@1 and save checkpoint
is_best = prec1 > best_prec1
best_prec1 = max(prec1, best_prec1)
save_checkpoint(
{
'epoch': epoch + 1,
'arch': configs.TRAIN.arch,
'state_dict': model.state_dict(),
'best_prec1': best_prec1,
'optimizer': optimizer.state_dict(),
}, is_best, os.path.join('trained_models', configs.output_name))
# Automatically perform PGD Attacks at the end of training
logger.info(pad_str(' Performing PGD Attacks '))
for pgd_param in configs.ADV.pgd_attack:
validate_pgd(val_loader, model, criterion, pgd_param[0], pgd_param[1],
configs, logger)
def cifar10_load_data(data_dir,
batch_size,
num_workers,
valid_size=0.1,
deterministic=False):
"""Load the CIFAR10 dataset.
The original training dataset is split into training and validation sets (code is
inspired by https://gist.github.com/kevinzakka/d33bf8d6c7f06a9d8c76d97a7879f5cb).
By default we use a 90:10 (45K:5K) training:validation split.
The output of torchvision datasets are PIL Image images of range [0, 1].
We transform them to Tensors of normalized range [-1, 1]
https://github.com/pytorch/tutorials/blob/master/beginner_source/blitz/cifar10_tutorial.py
Data augmentation: 4 pixels are padded on each side, and a 32x32 crop is randomly sampled
from the padded image or its horizontal flip.
This is similar to [1] and some other work that use CIFAR10.
[1] C.-Y. Lee, S. Xie, P. Gallagher, Z. Zhang, and Z. Tu. Deeply Supervised Nets.
arXiv:1409.5185, 2014
"""
transform = transforms.Compose([
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
transform_test = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
train_dataset = datasets.CIFAR10(root=data_dir,
train=True,
download=True,
transform=transform)
num_train = len(train_dataset)
indices = list(range(num_train))
split = int(np.floor(valid_size * num_train))
np.random.shuffle(indices)
train_idx, valid_idx = indices[split:], indices[:split]
train_sampler = SubsetRandomSampler(train_idx)
worker_init_fn = __deterministic_worker_init_fn if deterministic else None
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=batch_size,
sampler=train_sampler,
num_workers=num_workers,
pin_memory=True,
worker_init_fn=worker_init_fn)
valid_loader = None
if split > 0:
valid_sampler = SubsetRandomSampler(valid_idx)
valid_loader = torch.utils.data.DataLoader(
train_dataset,
batch_size=batch_size,
sampler=valid_sampler,
num_workers=num_workers,
pin_memory=True,
worker_init_fn=worker_init_fn)
testset = datasets.CIFAR10(root=data_dir,
train=False,
download=True,
transform=transform_test)
test_loader = torch.utils.data.DataLoader(testset,
batch_size=batch_size,
shuffle=False,
num_workers=num_workers,
pin_memory=True)
input_shape = __image_size(train_dataset)
# If validation split was 0 we use the test set as the validation set
return train_loader, valid_loader or test_loader, test_loader, input_shape
if not os.path.isdir(args.out_dir):
os.mkdir(args.out_dir)
random.seed(args.seed)
np.random.seed(args.seed) # to make node random permutation reproducible (not tested)
# Read image data using torchvision
is_train = args.split.lower() == 'train'
if args.dataset == 'mnist':
data = datasets.MNIST(args.data_dir, train=is_train, download=True)
assert args.compactness < 10, ('high compactness can result in bad superpixels on MNIST')
assert args.n_sp > 1 and args.n_sp < 28*28, (
'the number of superpixels cannot exceed the total number of pixels or be too small')
elif args.dataset == 'cifar10':
data = datasets.CIFAR10(args.data_dir, train=is_train, download=True)
assert args.compactness > 1, ('low compactness can result in bad superpixels on CIFAR-10')
assert args.n_sp > 1 and args.n_sp < 32*32, (
'the number of superpixels cannot exceed the total number of pixels or be too small')
else:
raise NotImplementedError('unsupported dataset: ' + args.dataset)
images = data.train_data if is_train else data.test_data
labels = data.train_labels if is_train else data.test_labels
if not isinstance(images, np.ndarray):
images = images.numpy()
if isinstance(labels, list):
labels = np.array(labels)
if not isinstance(labels, np.ndarray):
labels = labels.numpy()
generated1 = self.net1(x, self.mark1)
generated2 = self.net1(x, self.mark2)
generated = self.interp_z * generated1 + (1 - self.interp_z) * generated2
return self.net2(generated)
opt = TestOptions().parse()
train_transform = trn.Compose([trn.RandomHorizontalFlip(), trn.RandomCrop(32, padding=4),
trn.ToTensor(), trn.Normalize(
mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])])
test_transform = trn.Compose([trn.ToTensor(), trn.Normalize(
mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])])
if opt.dataset == 'cifar10':
train_data = dset.CIFAR10(opt.dataroot, train=True, transform=train_transform, download=True)
test_data = dset.CIFAR10(opt.dataroot, train=False, transform=test_transform)
num_classes = 10
else:
train_data = dset.CIFAR100(opt.dataroot, train=True, transform=train_transform, download=True)
test_data = dset.CIFAR100(opt.dataroot, train=False, transform=test_transform)
num_classes = 100
train_loader = torch.utils.data.DataLoader(
train_data, batch_size=opt.batch_size, shuffle=True,
num_workers=opt.prefetch, pin_memory=torch.cuda.is_available())
test_loader = torch.utils.data.DataLoader(
test_data, batch_size=opt.test_bs, shuffle=False,
num_workers=opt.prefetch, pin_memory=torch.cuda.is_available())
# Create model
def getData(name = 'cifar10', train_bs = 128, test_bs = 1000):
if name == 'svhn':
train_loader = torch.utils.data.DataLoader(
datasets.SVHN('../data', split = 'extra', download = True,
transform = transforms.Compose([
transforms.ToTensor()
])),
batch_size = train_bs, shuffle = True)
test_loader = torch.utils.data.DataLoader(
datasets.SVHN('../data', split = 'test', download = True, transform = transforms.Compose([
transforms.ToTensor()
])),
batch_size = test_bs, shuffle=False)
if name == 'mnist':
train_loader = torch.utils.data.DataLoader(
datasets.MNIST('../data', train = True, download = True,
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))
])),
batch_size = train_bs, shuffle = True)
test_loader = torch.utils.data.DataLoader(
datasets.MNIST('../data', train = False, transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))
])),
batch_size = test_bs, shuffle = False)
if name == 'cifar10':
transform_train = transforms.Compose([
#transforms.RandomCrop(32, padding = 4),
#transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
#transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
])
transform_test = transforms.Compose([
transforms.ToTensor(),
#transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
])
trainset = datasets.CIFAR10(root='../data', train = True, download = True, transform = transform_train)
train_loader = torch.utils.data.DataLoader(trainset, batch_size = train_bs, shuffle = True)
testset = datasets.CIFAR10(root='../data', train = False, download = True, transform = transform_test)
test_loader = torch.utils.data.DataLoader(testset, batch_size = test_bs, shuffle = False)
if name == 'cifar10da': # with data augumentation
transform_train = transforms.Compose([
transforms.RandomCrop(32, padding = 4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
])
transform_test = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
])
trainset = datasets.CIFAR10(root='../data', train = True, download = True, transform = transform_train)
train_loader = torch.utils.data.DataLoader(trainset, batch_size = train_bs, shuffle = True)
testset = datasets.CIFAR10(root='../data', train = False, download = True, transform = transform_test)
test_loader = torch.utils.data.DataLoader(testset, batch_size = test_bs, shuffle = False)
if name == 'cifar100':
transform_train = transforms.Compose([
transforms.RandomCrop(32, padding = 4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
])
transform_test = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
])
trainset = datasets.CIFAR100(root='../data', train = True, download = True, transform = transform_train)
train_loader = torch.utils.data.DataLoader(trainset, batch_size = train_bs, shuffle = True)
testset = datasets.CIFAR100(root='../data', train = False, download = True, transform = transform_test)
test_loader = torch.utils.data.DataLoader(testset, batch_size = test_bs, shuffle = False)
if name == 'tinyimagenet':
normalize = transforms.Normalize(mean = [0.44785526394844055, 0.41693055629730225, 0.36942949891090393],
std = [0.2928885519504547, 0.28230994939804077, 0.2889912724494934])
train_dataset = datasets.ImageFolder(
'../data/tiny-imagenet-200/train',
transforms.Compose([
transforms.RandomCrop(64, padding = 4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize,
]))
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size = train_bs, shuffle = True, num_workers = 4, pin_memory = False)
test_dataset = datasets.ImageFolder(
'../data/tiny-imagenet-200/val',
transforms.Compose([
transforms.ToTensor(),
normalize,
]))
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size = test_bs, shuffle = False)
return train_loader, test_loader
import torch.optim as optim
lr = 1e-4
cuda = False
G, D = Generator(), Discriminator()
if cuda:
G, D = G.cuda(), D.cuda()
G_optim, D_optim = optim.Adam(G.parameters(),
lr=lr), optim.Adam(D.parameters(), lr=lr)
# Load data
from torchvision import datasets, transforms
cifar10 = datasets.CIFAR10('data',
train=True,
download=True,
transform=transforms.ToTensor())
images = torch.stack([cifar10[i][0] for i in range(len(cifar10))])
orig_labels = torch.LongTensor([cifar10[i][1] for i in range(len(cifar10))])
labels = torch.zeros(images.size(0), 10).scatter_(1, orig_labels.view(-1, 1),
1)
if cuda:
images, orig_labels, labels = images.cuda(), orig_labels.cuda(
), labels.cuda()
# Load and save modal
def load_state():
if os.path.exists(prefix + '/checkpoint/G.data'):
G.load_state_dict(torch.load(prefix + '/checkpoint/G.data'))
if os.path.exists(prefix + '/checkpoint/D.data'):
def main():
if not torch.cuda.is_available():
logging.info('no gpu device available')
sys.exit(1)
np.random.seed(args.seed)
torch.cuda.set_device(args.gpu)
cudnn.benchmark = True
torch.manual_seed(args.seed)
cudnn.enabled = True
torch.cuda.manual_seed(args.seed)
logging.info('gpu device = %d' % args.gpu)
logging.info("args = %s", args)
genotype = eval("genotypes.%s" % args.arch)
logging.info('genotype = %s', genotype)
model = Network(CIFAR_CLASSES, genotype)
model = model.cuda()
test_epoch = 1
logging.info("param size = %fMB", utils.count_parameters_in_MB(model))
criterion = nn.CrossEntropyLoss()
criterion = criterion.cuda()
train_transform, test_transform = utils._data_transforms_cifar10(args)
test_data = dset.CIFAR10(root=args.data,
train=False,
download=True,
transform=test_transform)
test_queue = torch.utils.data.DataLoader(test_data,
batch_size=args.batch_size,
shuffle=False,
pin_memory=True,
num_workers=8)
logging.info("param size = %fMB", utils.count_parameters_in_MB(model))
optimizer = torch.optim.SGD(model.parameters(),
args.learning_rate,
momentum=args.momentum,
weight_decay=args.weight_decay)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
optimizer, float(args.epochs))
train_data = dset.CIFAR10(root=args.data,
train=True,
download=True,
transform=train_transform)
train_queue = torch.utils.data.DataLoader(train_data,
batch_size=args.batch_size,
pin_memory=True,
num_workers=8)
best_acc = 0.0
writer = SummaryWriter(args.save)
for epoch in range(args.epochs):
scheduler.step()
model.drop_path_prob = args.drop_path_prob * epoch / args.epochs
train_acc, train_obj = train(train_queue, model, criterion, optimizer)
logging.info('train_acc %f', train_acc)
valid_acc, valid_obj = infer(test_queue, model, criterion)
writer.add_scalar('trai_loss', train_obj)
writer.add_scalar('trian_acc', train_acc)
writer.add_scalar('val_loss', valid_obj)
writer.add_scalar('val_acc', valid_acc)
if valid_acc > best_acc:
best_acc = valid_acc
logging.info('epoch %d, valid_acc %f, best_acc %f', epoch,
valid_acc, best_acc)
utils.save(model, os.path.join(args.save, 'weights_retrain.pt'))
decay_equilibrium = args.decay_equilibrium
slurm = args.slurm
batchsize = args.batchsize
# TODO: add to argument parser
dataset_name = 'cifar10'
writer = SummaryWriter(comment="_CIFAR10_GAN")
net = VAEGAN(z_size=z_size, recon_level=recon_level).cuda()
# DATASET
if dataset_name == 'cifar10':
dataset = dset.CIFAR10(root=train_folder,
download=True,
transform=transforms.Compose([
transforms.Scale(z_size),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5),
(0.5, 0.5, 0.5)),
]))
dataloader = torch.utils.data.DataLoader(dataset,
batch_size=64,
shuffle=True,
num_workers=4)
# margin and equilibirum
margin = 0.35
equilibrium = 0.68
# mse_lambda = 1.0
# OPTIM-LOSS
# an optimizer for each of the sub-networks, so we can selectively backprop
def test(model):
kwargs = {'num_workers': 1, 'pin_memory': True} if args.cuda else {}
if args.dataset == 'cifar10':
test_loader = torch.utils.data.DataLoader(
datasets.CIFAR10('./data.cifar10',
train=False,
transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465),
(0.2023, 0.1994, 0.2010))
])),
batch_size=args.test_batch_size,
shuffle=False,
**kwargs)
elif args.dataset == 'cifar100':
test_loader = torch.utils.data.DataLoader(
datasets.CIFAR100('./data.cifar100',
train=False,
transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(
(0.4914, 0.4822, 0.4465),
(0.2023, 0.1994, 0.2010))
])),
batch_size=args.test_batch_size,
shuffle=False,
**kwargs)
elif args.dataset == 'imagenet':
# Data loading code
valdir = os.path.join(args.data, 'val')
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
test_loader = torch.utils.data.DataLoader(
datasets.ImageFolder(
valdir,
transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
normalize,
])),
batch_size=args.test_batch_size,
shuffle=False,
num_workers=16,
pin_memory=True)
else:
raise ValueError("No valid dataset is given.")
model.eval()
test_acc = 0
for data, target in test_loader:
if args.cuda:
data, target = data.cuda(), target.cuda()
data, target = Variable(data, volatile=True), Variable(target)
output = model(data)
pred = output.data.max(
1, keepdim=True)[1] # get the index of the max log-probability
prec1, prec5 = accuracy(output.data, target.data, topk=(1, 5))
test_acc += prec1.item()
print('\nTest set: Accuracy: {}/{} ({:.1f}%)\n'.format(
test_acc, len(test_loader), test_acc / len(test_loader)))
return np.round(test_acc / len(test_loader), 2)
def load_data(self, train_batch_size=64, test_batch_size=512):
"""
\Description : Load the dataset
\Args :
train_batch_size : size of the training batch size
test_batch_size : size og the test batch size
\Output :
train_loader : loader of the train batch
test_loader : loader of the test batch
"""
self.train_batch_size = train_batch_size
self.test_batch_size = test_batch_size
if self.dataset == "MNIST":
# Fetch training data
self.train_loader = torch.utils.data.DataLoader(
datasets.MNIST('data',
train=True,
download=True,
transform=transforms.Compose([
transforms.Resize((32, 32)),
transforms.ToTensor(),
transforms.Normalize((0.1307, ), (0.3081, ))
])),
batch_size=self.train_batch_size,
shuffle=True)
# Fetch test data
self.test_loader = torch.utils.data.DataLoader(
datasets.MNIST('data',
train=False,
transform=transforms.Compose([
transforms.Resize((32, 32)),
transforms.ToTensor(),
transforms.Normalize((0.1307, ), (0.3081, ))
])),
batch_size=self.test_batch_size,
shuffle=True)
elif self.dataset == "CIFAR10":
# Fetch training data
self.train_loader = torch.utils.data.DataLoader(
datasets.CIFAR10('data',
train=True,
download=True,
transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5),
(0.5, 0.5, 0.5))
])),
batch_size=self.train_batch_size,
shuffle=True)
# Fetch test data
self.test_loader = torch.utils.data.DataLoader(
datasets.CIFAR10('data',
train=False,
transform=transforms.Compose([
transforms.Resize((32, 32)),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5),
(0.5, 0.5, 0.5))
])),
batch_size=self.test_batch_size,
shuffle=True)
else:
raise ValueError(
"Invalid dataset name. Either choose 'MNIST' or 'CIFAR10'")
image_size = 32
net = Net()
print(net)
#hyperparametrers
learning_rate = 0.01
epochs = 15
logging_interval = 10
batch_size = 200
train_loader = torch.utils.data.DataLoader(
datasets.CIFAR10('../data_CIFAR', train=True, download=True,
transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))
])),
batch_size=batch_size, shuffle=True)
test_loader = torch.utils.data.DataLoader(
datasets.CIFAR10('../data_CIFAR', train=False, transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))
])),
batch_size=batch_size, shuffle=True)
#train the network
# create a stochastic gradient descent optimizer
optimizer = optim.SGD(net.parameters(),lr=learning_rate,momentum=0.9)
#create loss function - negative log likelihood loss
def load_dataset(dataset, path, noise=0, batch_size=100):
"""Load the specified dataset and form mini-batches of size batch_size.
Args
dataset (str) -- name of the dataset
path (str) -- path to the datasets
batch_size (int) -- mini-batch size
noise (float) -- percentage of noisy labels in the training set [0, 1.0]
"""
if not os.path.exists(path):
os.mkdir(path)
logger = logging.getLogger('train')
msg = 'Loading {}'.format(dataset)
msg = msg + ', corrupt labels with probability {}'.format(
noise) if noise > 0 else ''
logger.info(msg)
if dataset == 'mnist':
transform = transforms.Compose(
[transforms.ToTensor(),
transforms.Normalize((0.5, ), (1.0, ))])
# if not available locally, download mnist
if noise > 0:
train_set = MNISTNoisyLabels(noise=noise,
root=path,
train=True,
transform=transform,
download=False)
else:
train_set = datasets.MNIST(root=path,
train=True,
transform=transform,
download=True)
test_set = datasets.MNIST(root=path,
train=False,
transform=transform,
download=True)
train_loader = torch.utils.data.DataLoader(dataset=train_set,
batch_size=batch_size,
shuffle=True)
test_loader = torch.utils.data.DataLoader(dataset=test_set,
batch_size=batch_size,
shuffle=False)
elif dataset == 'cifar10':
data_augmentation = transforms.Compose([
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
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_test = transforms.Compose([
transforms.ToTensor(),
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]])
])
# if not available, download cifar10
if noise > 0: # no data augmentation on random labels
train_set = CIFAR10NoisyLabels(noise=noise,
root=path,
train=True,
transform=transform_test,
download=False)
else:
train_set = datasets.CIFAR10(root=path,
train=True,
download=True,
transform=data_augmentation)
train_loader = torch.utils.data.DataLoader(train_set,
batch_size=batch_size,
shuffle=True,
num_workers=4)
test_set = datasets.CIFAR10(root=path,
train=False,
download=True,
transform=transform_test)
test_loader = torch.utils.data.DataLoader(test_set,
batch_size=batch_size,
shuffle=False,
num_workers=4)
else:
raise ValueError("Unsupported dataset.")
return train_loader, test_loader
if __name__ == "__main__":
training_set_mnist = datasets.MNIST('./mnist_data', train=True, download=True,
transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))]))
train_loader_mnist = torch.utils.data.DataLoader(training_set_mnist, batch_size=128, shuffle=True)
test_loader_mnist = torch.utils.data.DataLoader(
datasets.MNIST('./mnist_data', train=False, transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))
])), batch_size=100, shuffle=True)
trainset_cifar10 = datasets.CIFAR10(root='./cifar10_data', train=True,
download=True, transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
]))
train_loader_cifar10 = torch.utils.data.DataLoader(trainset_cifar10, batch_size=128,
shuffle=True)
test_loader_cifar10 = torch.utils.data.DataLoader(
datasets.CIFAR10('./cifar10_data', train=False, transform=transforms.Compose([
transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])), batch_size=100, shuffle=True)
# load training and test set here:
training_set = datasets.CIFAR100(root='./cifar100_data', train=True,
download=True, transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
]))
train_loader = torch.utils.data.DataLoader(training_set, batch_size=128,
onnx.export(model, input_image, 'model.onnx')
# 您可以使用ONNX模型做很多事情，包括在不同平台和不同编程语言上运行推理。
# ## 代码实战
# In[18]:
import torch
from torch.utils.data import Dataset
from torchvision.transforms import ToTensor, Lambda
# In[16]:
train_data = datasets.CIFAR10(root="data",
train=True,
download=True,
transform=ToTensor())
# In[20]:
target_transform = Lambda(lambda y: torch.zeros(10, dtype=torch.float).
scatter_(dim=0, index=torch.tensor(y), value=1))
test_data = datasets.CIFAR10(root="data",
train=False,
download=True,
transform=target_transform)
# In[29]:
train_data.shape
torch.manual_seed(args.seed)
device = torch.device("cuda")
kwargs = {'batch_size': args.batch_size}
kwargs.update({'num_workers': 8, 'pin_memory': True, 'shuffle': True})
# Defining batch transforms
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
# Loading dataset
train_data = datasets.CIFAR10("~/Desktop/Datasets/",
train=True,
transform=transform)
test_data = datasets.CIFAR10("~/Desktop/Datasets/",
train=False,
transform=transform)
train_loader = torch.utils.data.DataLoader(train_data, **kwargs)
test_loader = torch.utils.data.DataLoader(test_data, **kwargs)
# Loading model
if args.arch == "my":
from Nets import *
model = Net().to(device)
def get_loaders_eval(dataset, args):
"""Get train and valid loaders for cifar10/tiny imagenet."""
if dataset == 'cifar10':
num_classes = 10
train_transform, valid_transform = _data_transforms_cifar10(args)
train_data = dset.CIFAR10(root=args.data,
train=True,
download=True,
transform=train_transform)
valid_data = dset.CIFAR10(root=args.data,
train=False,
download=True,
transform=valid_transform)
elif dataset == 'mnist':
num_classes = 10
train_transform, valid_transform = _data_transforms_mnist(args)
train_data = dset.MNIST(root=args.data,
train=True,
download=True,
transform=train_transform)
valid_data = dset.MNIST(root=args.data,
train=False,
download=True,
transform=valid_transform)
elif dataset.startswith('celeba'):
if dataset == 'celeba_64':
resize = 64
num_classes = 40
train_transform, valid_transform = _data_transforms_celeba64(
resize)
train_data = LMDBDataset(root=args.data,
name='celeba64',
train=True,
transform=train_transform,
is_encoded=True)
valid_data = LMDBDataset(root=args.data,
name='celeba64',
train=False,
transform=valid_transform,
is_encoded=True)
elif dataset in {'celeba_256'}:
num_classes = 1
resize = int(dataset.split('_')[1])
train_transform, valid_transform = _data_transforms_generic(resize)
train_data = LMDBDataset(root=args.data,
name='celeba',
train=True,
transform=train_transform)
valid_data = LMDBDataset(root=args.data,
name='celeba',
train=False,
transform=valid_transform)
else:
raise NotImplementedError
elif dataset.startswith('lsun'):
if dataset.startswith('lsun_bedroom'):
resize = int(dataset.split('_')[-1])
num_classes = 1
train_transform, valid_transform = _data_transforms_lsun(resize)
train_data = LSUN(root=args.data,
classes=['bedroom_train'],
transform=train_transform)
valid_data = LSUN(root=args.data,
classes=['bedroom_val'],
transform=valid_transform)
elif dataset.startswith('lsun_church'):
resize = int(dataset.split('_')[-1])
num_classes = 1
train_transform, valid_transform = _data_transforms_lsun(resize)
train_data = LSUN(root=args.data,
classes=['church_outdoor_train'],
transform=train_transform)
valid_data = LSUN(root=args.data,
classes=['church_outdoor_val'],
transform=valid_transform)
elif dataset.startswith('lsun_tower'):
resize = int(dataset.split('_')[-1])
num_classes = 1
train_transform, valid_transform = _data_transforms_lsun(resize)
train_data = LSUN(root=args.data,
classes=['tower_train'],
transform=train_transform)
valid_data = LSUN(root=args.data,
classes=['tower_val'],
transform=valid_transform)
else:
raise NotImplementedError
elif dataset.startswith('imagenet'):
num_classes = 1
resize = int(dataset.split('_')[1])
assert args.data.replace('/', '')[-3:] == dataset.replace(
'/', '')[-3:], 'the size should match'
train_transform, valid_transform = _data_transforms_generic(resize)
train_data = LMDBDataset(root=args.data,
name='imagenet-oord',
train=True,
transform=train_transform)
valid_data = LMDBDataset(root=args.data,
name='imagenet-oord',
train=False,
transform=valid_transform)
elif dataset.startswith('ffhq'):
num_classes = 1
resize = 256
train_transform, valid_transform = _data_transforms_generic(resize)
train_data = LMDBDataset(root=args.data,
name='ffhq',
train=True,
transform=train_transform)
valid_data = LMDBDataset(root=args.data,
name='ffhq',
train=False,
transform=valid_transform)
elif dataset == "audio":
num_classes = 1
train_data, valid_data, test_data = load(args=args)
else:
raise NotImplementedError
train_sampler, valid_sampler, test_sampler = None, None, None
if args.distributed:
train_sampler = torch.utils.data.distributed.DistributedSampler(
train_data)
valid_sampler = torch.utils.data.distributed.DistributedSampler(
valid_data)
test_sampler = torch.utils.data.distributed.DistributedSampler(
test_data)
train_queue = torch.utils.data.DataLoader(train_data,
batch_size=args.batch_size,
shuffle=(train_sampler is None),
sampler=train_sampler,
pin_memory=True,
num_workers=8,
drop_last=True)
valid_queue = torch.utils.data.DataLoader(valid_data,
batch_size=args.batch_size,
shuffle=(valid_sampler is None),
sampler=valid_sampler,
pin_memory=True,
num_workers=1,
drop_last=False)
test_queue = torch.utils.data.DataLoader(test_data,
batch_size=args.batch_size,
shuffle=(test_sampler is None),
sampler=test_sampler,
pin_memory=True,
num_workers=1,
drop_last=False)
return train_queue, valid_queue, num_classes, test_queue
from torchvision import datasets, transforms
from torch.optim.lr_scheduler import MultiStepLR
from nni.compression.pytorch import ModelSpeedup
from examples.model_compress.models.cifar10.vgg import VGG
from nni.compression.pytorch.utils.counter import count_flops_params
from nni.algorithms.compression.v2.pytorch.pruning.basic_pruner import L1NormPruner, L2NormPruner
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
normalize = transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
g_epoch = 0
train_loader = torch.utils.data.DataLoader(
datasets.CIFAR10('./data', train=True, transform=transforms.Compose([
transforms.RandomHorizontalFlip(),
transforms.RandomCrop(32, 4),
transforms.ToTensor(),
normalize,
]), download=True),
batch_size=128, shuffle=True)
test_loader = torch.utils.data.DataLoader(
datasets.CIFAR10('./data', train=False, transform=transforms.Compose([
transforms.ToTensor(),
normalize,
])),
batch_size=128, shuffle=False)
def trainer(model, optimizer, criterion):
global g_epoch
model.train()
for batch_idx, (data, target) in enumerate(train_loader):
def main():
if not torch.cuda.is_available():
logging.info('no gpu device available')
sys.exit(1)
np.random.seed(args.seed)
torch.cuda.set_device(args.gpu)
device = torch.device('cuda')
cudnn.benchmark = True
torch.manual_seed(args.seed)
cudnn.enabled = True
torch.cuda.manual_seed(args.seed)
logging.info('gpu device = %d' % args.gpu)
logging.info("args = %s", args)
criterion = nn.CrossEntropyLoss()
criterion = criterion.to(device)
model = Network(args.init_channels, CIFAR_CLASSES, args.layers, criterion)
model = model.to(device)
logging.info("param size = %fMB", utils.count_parameters_in_MB(model))
optimizer = torch.optim.SGD(model.parameters(),
args.learning_rate,
momentum=args.momentum,
weight_decay=args.weight_decay)
train_transform, valid_transform = utils._data_transforms_cifar10(args)
train_data = dset.CIFAR10(root=args.data,
train=True,
download=True,
transform=train_transform)
num_train = len(train_data)
indices = list(range(num_train))
split = int(np.floor(args.train_portion * num_train))
train_queue = torch.utils.data.DataLoader(
train_data,
batch_size=args.batch_size,
sampler=torch.utils.data.sampler.SubsetRandomSampler(indices[:split]),
pin_memory=True,
num_workers=0)
valid_queue = torch.utils.data.DataLoader(
train_data,
batch_size=32,
sampler=torch.utils.data.sampler.SubsetRandomSampler(
indices[split:num_train]),
pin_memory=True,
num_workers=0)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
optimizer, float(args.epochs), eta_min=args.learning_rate_min)
architect = Architect(model, args)
for epoch in range(args.epochs):
scheduler.step()
lr = scheduler.get_lr()[0]
logging.info('epoch %d lr %e', epoch, lr)
genotype = model.genotype()
logging.info('genotype = %s', genotype)
print(F.softmax(model.alphas_normal, dim=-1))
print(F.softmax(model.alphas_reduce, dim=-1))
# training
train_acc, train_obj = train(train_queue, valid_queue, model,
architect, criterion, optimizer, lr)
logging.info('train_acc %f', train_acc)
# validation
valid_acc, valid_obj = infer(valid_queue, model, criterion)
logging.info('valid_acc %f', valid_acc)
utils.save(model, os.path.join(args.save, 'weights.pt'))
def get_dataset(args):
""" return given network
"""
if args.dataset == 'mnist':
train_dataset = datasets.MNIST(MNIST_PATH,
download=True,
transform=transforms.Compose([
transforms.Resize((32, 32)),
transforms.ToTensor(),
]))
val_dataset = datasets.MNIST(MNIST_PATH,
train=False,
download=True,
transform=transforms.Compose([
transforms.Resize((32, 32)),
transforms.ToTensor(),
]))
elif args.dataset == 'cifar10':
train_dataset = datasets.CIFAR10(
CIFAR10_PATH,
download=True,
transform=transforms.Compose([
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.RandomRotation(15),
transforms.ToTensor(),
transforms.Normalize(mean=[0.49139968, 0.48215827, 0.44653124],
std=[0.24703233, 0.24348505, 0.26158768]),
]))
val_dataset = datasets.CIFAR10(
CIFAR10_PATH,
train=False,
download=True,
transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(mean=[0.49139968, 0.48215827, 0.44653124],
std=[0.24703233, 0.24348505, 0.26158768]),
]))
elif args.dataset == 'cifar100':
train_dataset = datasets.CIFAR100(
CIFAR100_PATH,
download=True,
transform=transforms.Compose([
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.RandomRotation(15),
transforms.ToTensor(),
transforms.Normalize(mean=[
0.5070751592371323, 0.48654887331495095, 0.4409178433670343
],
std=[
0.2673342858792401,
0.2564384629170883,
0.27615047132568404
]),
]))
val_dataset = datasets.CIFAR100(
CIFAR100_PATH,
train=False,
download=True,
transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(mean=[
0.5070751592371323, 0.48654887331495095, 0.4409178433670343
],
std=[
0.2673342858792401,
0.2564384629170883,
0.27615047132568404
]),
]))
elif args.dataset == 'imagenet':
input_image_size = 224
scale = 256 / 224
train_dataset_path = os.path.join(ILSVRC2012_path, 'train')
val_dataset_path = os.path.join(ILSVRC2012_path, 'val')
train_dataset = datasets.ImageFolder(
train_dataset_path,
transforms.Compose([
transforms.RandomResizedCrop(input_image_size),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]),
]))
val_dataset = datasets.ImageFolder(
val_dataset_path,
transforms.Compose([
transforms.Resize(int(input_image_size * scale)),
transforms.CenterCrop(input_image_size),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]),
]))
elif args.dataset == 'imagenet_32_noise':
input_image_size = 224
scale = 256 / 224
train_dataset_path = os.path.join(IMAGENET_32_NOISE_PATH, 'train')
val_dataset_path = os.path.join(IMAGENET_32_NOISE_PATH, 'val')
train_dataset = datasets.ImageFolder(
train_dataset_path,
transforms.Compose([
transforms.RandomResizedCrop(input_image_size),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]),
]))
val_dataset = datasets.ImageFolder(
val_dataset_path,
transforms.Compose([
transforms.Resize(int(input_image_size * scale)),
transforms.CenterCrop(input_image_size),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]),
]))
elif args.dataset == 'imagenet_32_noise_2':
input_image_size = 224
scale = 256 / 224
train_dataset_path = os.path.join(IMAGENET_32_NOISE_PATH_2, 'train')
val_dataset_path = os.path.join(IMAGENET_32_NOISE_PATH_2, 'val')
train_dataset = datasets.ImageFolder(
train_dataset_path,
transforms.Compose([
transforms.RandomResizedCrop(input_image_size),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]),
]))
val_dataset = datasets.ImageFolder(
val_dataset_path,
transforms.Compose([
transforms.Resize(int(input_image_size * scale)),
transforms.CenterCrop(input_image_size),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]),
]))
else:
print('the dataset name you have entered is not supported yet')
sys.exit()
return train_dataset, val_dataset
def load_data(self, img_size=32):
data_dir = '/home/zengyuyuan/data/CIFAR10'
data_transforms = {
'train':
transforms.Compose([
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize([0.4914, 0.4822, 0.4465],
[0.2023, 0.1994, 0.2010])
]),
'val':
transforms.Compose([
transforms.ToTensor(),
transforms.Normalize([0.4914, 0.4822, 0.4465],
[0.2023, 0.1994, 0.2010])
])
}
if self.dataset == 'cifar-10':
data_train = datasets.CIFAR10(root=data_dir,
transform=data_transforms['train'],
train=True,
download=True)
data_test = datasets.CIFAR10(root=data_dir,
transform=data_transforms['val'],
train=False,
download=True)
if self.dataset == 'cifar-100':
data_train = datasets.CIFAR100(root=data_dir,
transform=data_transforms['train'],
train=True,
download=True)
data_test = datasets.CIFAR100(root=data_dir,
transform=data_transforms['val'],
train=False,
download=True)
if self.dataset == 'mnist':
data_dir = '/home/zengyuyuan/data/MNIST'
mnist_transforms = {
'train':
transforms.Compose([
transforms.ToTensor(),
transforms.Normalize([
0.1307,
], [0.3081])
]),
'val':
transforms.Compose([
transforms.ToTensor(),
transforms.Normalize([
0.1307,
], [0.3081])
])
}
data_train = datasets.MNIST(root=data_dir,
transform=mnist_transforms['train'],
train=True,
download=True)
data_test = datasets.MNIST(root=data_dir,
transform=mnist_transforms['val'],
train=False,
download=True)
if self.dataset == 'VOCpart':
data_train = VOCPart(VOC_path,
train=True,
requires=['img'],
size=img_size)
data_test = VOCPart(VOC_path,
train=False,
requires=['img'],
size=img_size)
image_datasets = {'train': data_train, 'val': data_test}
# change list to Tensor as the input of the models
dataloaders = {}
dataloaders['train'] = torch.utils.data.DataLoader(
image_datasets['train'],
batch_size=self.batch_size,
pin_memory=True,
shuffle=True,
worker_init_fn=set_work_init_fn(self.seed),
num_workers=16)
dataloaders['val'] = torch.utils.data.DataLoader(
image_datasets['val'],
batch_size=self.batch_size,
pin_memory=True,
shuffle=False,
worker_init_fn=set_work_init_fn(self.seed),
num_workers=16)
dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'val']}
return dataloaders, dataset_sizes
for dt, _ in it.islice(train_loader, 0, 1, 1):
pass
init_loader = torch.utils.data.DataLoader(dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=5,
drop_last=True,
pin_memory=True)
model = Glow_((args.batch_size, 3, 64, 64), args).cuda()
elif args.data == 'cifar10':
train_loader = torch.utils.data.DataLoader(datasets.CIFAR10(args.data_dir,
train=True,
download=True,
transform=tf),
batch_size=args.batch_size,
shuffle=True,
num_workers=10,
drop_last=True)
test_loader = torch.utils.data.DataLoader(datasets.CIFAR10(args.data_dir,
train=False,
transform=tf),
batch_size=args.batch_size,
shuffle=False,
num_workers=10)
# data dependant init
init_loader = torch.utils.data.DataLoader(datasets.CIFAR10(args.data_dir,