scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer,
T_0=3)
################################################
# Dataset and train-test helpers
################################################
transform_val = transforms.Compose([
transforms.Resize(224),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
data_root = args.IMAGENET_DIR
train_ds = ImageNet(data_root,
split='train',
download=False,
transform=transform_val)
train_loader = torch.utils.data.DataLoader(train_ds,
batch_size=args.BATCHSIZE,
shuffle=True,
drop_last=False,
num_workers=args.NUM_WORKERS,
pin_memory=True)
val_ds = ImageNet(data_root,
split='val',
download=False,
transform=transform_val)
val_loader = torch.utils.data.DataLoader(val_ds,
batch_size=args.BATCHSIZE,
shuffle=True,
args = parse_args()
# Location of test images
images_fpaths = get_image_paths()
if args.image_name is not None:
idx = [
i for i in range(len(images_fpaths))
if args.image_name in images_fpaths[i]
][0]
images_fpaths = [images_fpaths[idx]]
# Get the actual labels from the filenames
labels_act = parse_labels_from_paths(images_fpaths)
# Get ImageNet dataset for interpreting the outputs of the network
imagenet = ImageNet(IMAGENET_LOCATION, split='val', download=True)
# For each input image, transform to be same size and orientation
transform = init_image_transforms(args.image_size)
# Put all images into a big tensor for putting into model
imgs = read_images_as_tensors(images_fpaths, transform)
# Instantiate model
model = TransparentSqueezeNet(pretrained=True)
# Pass images through model
output_opaque = model(imgs)
# Predicted label
labels_pred = to_readable(output_opaque, imagenet)
dtype = args.dtype
num_workers = args.num_workers
batch_size = args.batch_size * len(device_ids)
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
val_transform = transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
normalize,
])
if not args.use_lmdb:
val_set = ImageNet(args.data_path, split='val', transform=val_transform)
else:
val_set = ImageLMDB(os.path.join(args.data_path, 'val.lmdb'),
transform=val_transform)
val_data = DataLoader(val_set,
batch_size,
False,
pin_memory=True,
num_workers=num_workers,
drop_last=False)
model_setting = set_model(0, args.norm_layer, args.activation)
try:
model = get_model(args.model, alpha=args.alpha, **model_setting)
def load_dataset(dataset_name, **kwargs):
""" Loads the specified dataset and returns a PyTorch dataset object.
Applies the standard transformations for said dataset by default.
"""
data_path = pathlib.Path('data').resolve()
if dataset_name == 'cifar10':
from torchvision.datasets import CIFAR10
# This is the standard normalization transformation
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465),
(0.2023, 0.1994, 0.2010)),
])
# User can specify to load the training set; loads the test set by default.
train = kwargs.pop('train', False)
dataset = CIFAR10(data_path,
train=train,
transform=transform,
download=True)
elif dataset_name == 'cifar100':
from torchvision.datasets import CIFAR100
# This is the standard normalization transformation
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465),
(0.2023, 0.1994, 0.2010)),
])
# User can specify to load the training set; loads the test set by default.
train = kwargs.pop('train', False)
dataset = CIFAR100(data_path,
train=train,
transform=transform,
download=True)
elif dataset_name == 'imagenet':
# Requires imagenet to be downloaded locally
from torchvision.datasets import ImageNet
# Standard transformation
transform = transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
])
dataset = ImageNet(data_path / 'imagenet',
split='val',
transform=transform)
elif dataset_name == 'cifar10r':
from data.nonstationary_datasets import CIFAR10R
dataset = CIFAR10R()
elif dataset_name == 'cifar100r':
from data.nonstationary_datasets import CIFAR100R
dataset = CIFAR100R()
elif dataset_name == 'cifar10gb':
from data.nonstationary_datasets import CIFAR10GB
dataset = CIFAR10GB()
elif dataset_name == 'cifar100gb':
from data.nonstationary_datasets import CIFAR100GB
dataset = CIFAR100GB()
elif dataset_name == 'cifar10imba':
from data.imbalanced_datasets import CIFAR10Imba
dataset = CIFAR10Imba(class_ratios=kwargs['class_ratios'])
else:
raise NotImplementedError
return dataset
def main(argv=None):
path = '/home/alta/uncertainty/am969/pytorch/datasets/imagenet/tars'
dataset1 = ImageNet(root=path, split='train', download=True)
dataset2 = ImageNet(root=path, split='val ', download=True)
#input_transform = transforms.Compose([
# transforms.Resize(224,PIL.Image.BICUBIC), transforms.ToTensor(),
# transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
#])
input_transform = transforms.Compose([
transforms.Resize(256, PIL.Image.BICUBIC),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
test_dataset = ImageNet(
DATA_ROOT,
split="val",
transform=input_transform,
target_transform=None,
download=True,
)
test_loader = DataLoader(
test_dataset,
batch_size=64,
shuffle=False,
num_workers=4,
pin_memory=True,
)
model = model.cuda()
model.eval()
def main_worker(gpu, ngpus_per_node, args):
global best_acc1
args.gpu = gpu
if args.gpu is not None:
print("Use GPU: {} for training".format(args.gpu))
if args.distributed:
if args.dist_url == "env://" and args.rank == -1:
args.rank = int(os.environ["RANK"])
if args.multiprocessing_distributed:
# For multiprocessing distributed training, rank needs to be the
# global rank among all the processes
args.rank = args.rank * ngpus_per_node + gpu
dist.init_process_group(backend=args.dist_backend,
init_method=args.dist_url,
world_size=args.world_size,
rank=args.rank)
# 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 = getattr(resnet, args.arch)()
if not torch.cuda.is_available():
print('using CPU, this will be slow')
elif args.distributed:
# For multiprocessing distributed, DistributedDataParallel constructor
# should always set the single device scope, otherwise,
# DistributedDataParallel will use all available devices.
if args.gpu is not None:
torch.cuda.set_device(args.gpu)
model.cuda(args.gpu)
# When using a single GPU per process and per
# DistributedDataParallel, we need to divide the batch size
# ourselves based on the total number of GPUs we have
args.batch_size = int(args.batch_size / ngpus_per_node)
args.workers = int(
(args.workers + ngpus_per_node - 1) / ngpus_per_node)
model = torch.nn.parallel.DistributedDataParallel(
model, device_ids=[args.gpu])
else:
model.cuda()
# DistributedDataParallel will divide and allocate batch_size to all
# available GPUs if device_ids are not set
model = torch.nn.parallel.DistributedDataParallel(model)
elif args.gpu is not None:
torch.cuda.set_device(args.gpu)
model = model.cuda(args.gpu)
else:
# DataParallel will divide and allocate batch_size to all available GPUs
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()
# define loss function (criterion) and optimizer
criterion = nn.CrossEntropyLoss().cuda(args.gpu)
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))
if args.gpu is None:
checkpoint = torch.load(args.resume)
else:
# Map model to be loaded to specified single gpu.
loc = 'cuda:{}'.format(args.gpu)
checkpoint = torch.load(args.resume, map_location=loc)
args.start_epoch = checkpoint['epoch']
best_acc1 = checkpoint['best_acc1']
if args.gpu is not None:
# best_acc1 may be from a checkpoint from a different GPU
best_acc1 = best_acc1.to(args.gpu)
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
img_net = ImageNet(args.data)
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(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.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, args)
return
for epoch in range(args.start_epoch, args.epochs):
if args.distributed:
train_sampler.set_epoch(epoch)
adjust_learning_rate(optimizer, epoch, args)
# train for one epoch
train(train_loader, model, criterion, optimizer, epoch, args)
# evaluate on validation set
acc1 = validate(val_loader, model, criterion, args)
# remember best acc@1 and save checkpoint
is_best = acc1 > best_acc1
best_acc1 = max(acc1, best_acc1)
if not args.multiprocessing_distributed or (
args.multiprocessing_distributed
and args.rank % ngpus_per_node == 0):
save_checkpoint(
{
'epoch': epoch + 1,
'arch': args.arch,
'state_dict': model.state_dict(),
'best_acc1': best_acc1,
'optimizer': optimizer.state_dict(),
}, is_best)
# -*- coding: utf-8 -*- # @Author : DevinYang([email protected]) import argparse from torchvision.datasets import ImageNet from torchtoolbox.tools import check_dir from torchtoolbox.tools.convert_lmdb import generate_lmdb_dataset, raw_reader parser = argparse.ArgumentParser(description='Convert a ImageFolder dataset to LMDB format.') parser.add_argument('--data-dir', type=str, required=True, help='ImageFolder path, this param will give to ImageFolder Dataset.') parser.add_argument('--save-dir', type=str, required=True, help='Save dir.') parser.add_argument('--download', action='store_true', help='download dataset.') parser.add_argument('-j', dest='num_workers', type=int, default=0) parser.add_argument('--write-frequency', type=int, default=5000) parser.add_argument('--max-size', type=float, default=1.0, help='Maximum size database, this is rate, default is 1T, final setting would be ' '1T * `this param`') args = parser.parse_args() check_dir(args.save_dir) train_data_set = ImageNet(args.data_dir, 'train', args.download, loader=raw_reader) val_data_set = ImageNet(args.data_dir, 'val', args.download, loader=raw_reader) if __name__ == '__main__': generate_lmdb_dataset(train_data_set, args.save_dir, 'train', args.num_workers, args.max_size, args.write_frequency) generate_lmdb_dataset(val_data_set, args.save_dir, 'val', args.num_workers, args.max_size, args.write_frequency)
def experiment(num_shared_classes, percent_shared_data, n_epochs=200,batch_size=128, eps=.3, adv_steps=100, learning_rate=.0004, gpu_num=1,adv_training=False,task="CIFAR100"):
print("epochs,batch_size,eps,adv_steps,learning_rate,task")
print(n_epochs,batch_size,eps,adv_steps,learning_rate,task)
cuda = torch.cuda.is_available()
transform_test = transforms.Compose(
[transforms.ToTensor(),transforms.Normalize((0.5070751592371323, 0.48654887331495095, 0.4409178433670343), (0.2673342858792401, 0.2564384629170883, 0.27615047132568404))])
transform_train = transforms.Compose([
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.5070751592371323, 0.48654887331495095, 0.4409178433670343), (0.2673342858792401, 0.2564384629170883, 0.27615047132568404)),
])
if task.upper() == "CIFAR100":
transform_test = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
])
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))
])
train_data = CIFAR100("data/",transform=transform_train, download=False)
test_data = CIFAR100("data/", train=False, transform=transform_test, download=False)
elif task.upper() == "IMAGENET":
train_data = ImageNet('data/imagenet', split='train', download=False)
test_data = ImageNet('data/imagenet', split='val', download=False)
elif task.upper() == "FASHIONMNIST":
transform = transforms.Compose([transforms.Lambda(lambda image: image.convert('RGB')),
transforms.ToTensor()
])
train_data = FashionMNIST('data/fashionmnist',transform=transform, train=True, download=False)
test_data = FashionMNIST('data/fashionmnist', transform=transform, train=False, download=False)
else:
train_data = CIFAR10("data/",transform=transform_train,download=False)
test_data = CIFAR10("data/", train=False, transform=transform_test,download=False)
# model1 = ResNet(ResidualBlock, [2, 2, 2],num_classes=10)
# model2 = ResNet(ResidualBlock, [2, 2, 2],num_classes=10)
all_classes = set([x[1] for x in train_data])
shared_classes = random.sample(all_classes, num_shared_classes)
split_classes = [c for c in all_classes if c not in shared_classes] # get classes not shared
if len(split_classes) % 2 == 1: # if we have an odd #, randomly remove one so that number of classes will be the same for each model
split_classes.pop(random.randint(0, len(split_classes) - 1))
model1_split = random.sample(split_classes, len(split_classes) // 2)
model2_split = [c for c in split_classes if c not in model1_split]
model1_classes = model1_split
model2_classes = model2_split
model1_classes.sort()
model2_classes.sort()
# DEBUG:
print("shared classes: {}".format(shared_classes))
print("model1 classes: {}".format(model1_classes))
print("model2 classes: {}".format(model2_classes))
model1_x_train = []
model1_y_train = []
model2_x_train = []
model2_y_train = []
shared_x_train = []
shared_y_train = []
# train data splits
for index in range(len(train_data)):
current_class = train_data[index][1]
# model 1
if current_class in model1_classes:
model1_x_train.append(train_data[index][0])
model1_y_train.append(train_data[index][1])
# model 2
if current_class in model2_classes:
model2_x_train.append(train_data[index][0])
model2_y_train.append(train_data[index][1])
# split by percentage for classes per model1
if percent_shared_data < 100:
new_model1_x_train = []
new_model1_y_train = []
for curr_class in model1_classes:
temp_data_x = []
temp_data_y = []
# get all examples of class
for i in range(len(model1_x_train)):
if(model1_y_train[i] == curr_class):
temp_data_x.append(model1_x_train[i])
temp_data_y.append(model1_y_train[i])
# split data by half the size
total_size = len(temp_data_x)
shared_size = int(total_size * .5)
shared_indices = random.sample(list(range(len(temp_data_x))),shared_size)
new_model1_x_train += [temp_data_x[i] for i in shared_indices]
new_model1_y_train += [temp_data_y[i] for i in shared_indices]
# split for model2
new_model2_x_train = []
new_model2_y_train = []
for curr_class in model2_classes:
temp_data_x = []
temp_data_y = []
# get all examples of class
for i in range(len(model2_x_train)):
if(model2_y_train[i] == curr_class):
temp_data_x.append(model2_x_train[i])
temp_data_y.append(model2_y_train[i])
# split data by half the size
total_size = len(temp_data_x)
shared_size = int(total_size * .5)
shared_indices = random.sample(list(range(len(temp_data_x))),shared_size)
new_model2_x_train += [temp_data_x[i] for i in shared_indices]
new_model2_y_train += [temp_data_y[i] for i in shared_indices]
# rewrite dataset
model1_x_train = new_model1_x_train
model1_y_train = new_model1_y_train
model2_x_train = new_model2_x_train
model2_y_train = new_model2_y_train
# Carry out datasplitting for shared classes and add to datasets
for shared_class in shared_classes:
all_examples_x_train = []
all_examples_y_train = []
# get all examples of class
for index in range(len(train_data)):
current_class = train_data[index][1]
if current_class == shared_class:
all_examples_x_train.append(train_data[index][0])
all_examples_y_train.append(train_data[index][1])
# find max number of samples per model (set to be amount of examples if data is completely disjoint)
max_examples = len(all_examples_x_train) // 2
# get shared examples
shared_examples_x_train = []
shared_examples_y_train = []
num_shared_examples = max_examples * percent_shared_data // 100
for _ in range(num_shared_examples):
random_int = random.randint(0, len(all_examples_x_train) - 1)
shared_examples_x_train.append(all_examples_x_train.pop(random_int))
shared_examples_y_train.append(all_examples_y_train.pop(random_int))
# get disjoint examples
disjoint_examples = max_examples - len(shared_examples_x_train)
model1_examples_x_train = []
model1_examples_y_train = []
model2_examples_x_train = []
model2_examples_y_train = []
for _ in range(disjoint_examples):
model1_rand_int = random.randint(0, len(all_examples_x_train) - 1)
model1_examples_x_train.append(all_examples_x_train.pop(model1_rand_int))
model1_examples_y_train.append(all_examples_y_train.pop(model1_rand_int))
model2_rand_int = random.randint(0, len(all_examples_x_train) - 1)
model2_examples_x_train.append(all_examples_x_train.pop(model2_rand_int))
model2_examples_y_train.append(all_examples_y_train.pop(model2_rand_int))
# add to the datasets for the model
model1_x_train = shared_examples_x_train + model1_x_train + model1_examples_x_train
model1_y_train = shared_examples_y_train + model1_y_train + model1_examples_y_train
model2_x_train = shared_examples_x_train + model2_x_train + model2_examples_x_train
model2_y_train = shared_examples_y_train + model2_y_train + model2_examples_y_train
#print(model1_y_train)
# assign mapping for new classes
model1_class_mapping = {}
model2_class_mapping = {}
model1_classes_inc = 0
# go through model1 and assign unique classes to incrimental int starting at 0
for index in range(len(model1_y_train)):
# if it doesn't exist assign
if model1_y_train[index] not in model1_class_mapping.keys():
model1_class_mapping[model1_y_train[index]] = model1_classes_inc
model1_classes_inc += 1
# append assigned token
model1_y_train[index] = model1_class_mapping[model1_y_train[index]]
model2_classes_inc = 0
# go through model2 and assign unique classes to incrimental int starting at 0
for index in range(len(model2_y_train)):
# if it doesn't exist in model2 OR in model1, assign it
if model2_y_train[index] not in model2_class_mapping.keys() and model2_y_train[index] not in model1_class_mapping.keys():
model2_class_mapping[model2_y_train[index]] = model2_classes_inc
model2_y_train[index] = model2_classes_inc
model2_classes_inc += 1
elif model2_y_train[index] in model1_class_mapping.keys():
model2_y_train[index] = model1_class_mapping[model2_y_train[index]]
else:
model2_y_train[index] = model2_class_mapping[model2_y_train[index]]
model1_x_test = []
model1_y_test = []
model2_x_test = []
model2_y_test = []
shared_x_test = []
shared_y_test = []
# test data splits
for index in range(len(test_data)):
current_class = test_data[index][1]
# model 1
if current_class in model1_classes:
model1_x_test.append(test_data[index][0])
model1_y_test.append(test_data[index][1])
# model 2
if current_class in model2_classes:
model2_x_test.append(test_data[index][0])
model2_y_test.append(test_data[index][1])
# shared classes for eval
if current_class in shared_classes:
shared_x_test.append(test_data[index][0])
shared_y_test.append(test_data[index][1])
model1_x_test += shared_x_test
model1_y_test += shared_y_test
model2_x_test += shared_x_test
model2_y_test += shared_y_test
for index in range(len(model1_y_test)):
model1_y_test[index] = model1_class_mapping[model1_y_test[index]]
for index in range(len(model2_y_test)):
if model2_y_test[index] in model1_class_mapping.keys():
model2_y_test[index] = model1_class_mapping[model2_y_test[index]]
else:
model2_y_test[index] = model2_class_mapping[model2_y_test[index]]
model1_classes_len= len(set([item for item in model1_y_train]))
model2_classes_len = len(set([item for item in model2_y_train]))
if task.upper() == "CIFAR100":
model1 = models.wide_resnet50_2()
model2 = models.wide_resnet50_2()
#
model1.fc = nn.Linear(2048, model1_classes_len)
model2.fc = nn.Linear(2048, model2_classes_len)
elif task.upper() == "IMAGENET":
model1 = models.wide_resnet50_2()
model2 = models.wide_resnet50_2()
model1.fc = nn.Linear(2048, model1_classes_len)
model2.fc = nn.Linear(2048, model2_classes_len)
elif task.upper() == "FASHIONMNIST":
model1 = models.resnet18()
model2 = models.resnet18()
model1.fc = nn.Linear(512, model1_classes_len)
model2.fc = nn.Linear(512, model2_classes_len)
else:
# Get model (using ResNet50 for now)
model1 = models.resnet50()
model2 = models.resnet50()
model1.fc = nn.Linear(2048, model1_classes_len)
model2.fc = nn.Linear(2048, model2_classes_len)
cuda = torch.cuda.is_available()
if gpu_num in range(torch.cuda.device_count()):
device = torch.device('cuda:'+str(gpu_num) if cuda else 'cpu')
torch.cuda.set_device(device)
else:
device = torch.device('cpu')
# Model Training
model1 = model1.to(device)
model2 = model2.to(device)
criterion1 = nn.CrossEntropyLoss()
optimizer1 = optim.AdamW(model1.parameters(), lr=learning_rate)
scheduler1 = optim.lr_scheduler.MultiStepLR(optimizer1,milestones=[60, 120, 160], gamma=.2) #learning rate decay
criterion2 = nn.CrossEntropyLoss()
optimizer2 = optim.AdamW(model2.parameters(), lr=learning_rate)
scheduler2 = optim.lr_scheduler.MultiStepLR(optimizer2,milestones=[60, 120, 160], gamma=.2) #learning rate decay
# zip together two lists
train_set1 = list(zip(model1_x_train, model1_y_train))
# create trainloader 1
trainloader_1 = torch.utils.data.DataLoader(train_set1, batch_size=batch_size,
shuffle=True, num_workers=2)
# create trainloader 2
# zip together two lists
train_set2 = list(zip(model2_x_train, model2_y_train))
# create trainloader 1
trainloader_2 = torch.utils.data.DataLoader(train_set2, batch_size=batch_size,
shuffle=True, num_workers=2)
# TODO change this
num_adv_batchs = 2 if adv_training else 0
adv_batches = random.sample(range(len(trainloader_1)), num_adv_batchs)
#print("adv_batches:", adv_batches)
# train model 1
for epoch in tqdm(range(n_epochs),desc="Epoch"): # loop over the dataset multiple times
running_loss = 0.0
for i, data in enumerate(trainloader_1, 0):
if cuda:
data = tuple(d.cuda() for d in data)
# get the inputs; data is a list of [inputs, labels]
inputs, labels = data
# zero the parameter gradients
optimizer1.zero_grad()
# forward + backward + optimize
# train adversarial
# if i in adv_batches:
# print("adv training!")
# adversary = LinfPGDAttack(
# model1, loss_fn=nn.CrossEntropyLoss(reduction="sum"), eps=eps,
# nb_iter=adv_steps, eps_iter=0.01, rand_init=True, clip_min=0.0, clip_max=1.0,
# targeted=False)
# inputs = adversary.perturb(inputs, labels)
outputs = model1(inputs)
loss = criterion1(outputs, labels)
loss.backward()
optimizer1.step()
# print statistics
running_loss += loss.item()
if i % 2000 == 1999: # print every 2000 mini-batches
print('[%d, %5d] loss: %.3f' %
(epoch + 1, i + 1, running_loss / 2000))
running_loss = 0.0
print('Finished Training model1')
# train model 2
for epoch in tqdm(range(n_epochs),desc="Epoch"): # loop over the dataset multiple times
running_loss = 0.0
for i, data in enumerate(trainloader_2, 0):
if cuda:
data = tuple(d.cuda() for d in data)
# get the inputs; data is a list of [inputs, labels]
inputs, labels = data
# zero the parameter gradients
optimizer2.zero_grad()
# forward + backward + optimize
outputs = model2(inputs)
loss = criterion2(outputs, labels)
loss.backward()
optimizer2.step()
# print statistics
running_loss += loss.item()
if i % 2000 == 1999: # print every 2000 mini-batches
print('[%d, %5d] loss: %.3f' %
(epoch + 1, i + 1, running_loss / 2000))
running_loss = 0.0
print('Finished Training model2')
model1 = model1.to("cpu")
model2 = model2.to("cpu")
# convert shared classes to new labels
for index in range(len(shared_y_test)):
if shared_y_test[index] in model1_class_mapping.keys():
shared_y_test[index] = model1_class_mapping[shared_y_test[index]]
else:
shared_y_test[index] = model2_class_mapping[shared_y_test[index]]
shared_y_test = torch.Tensor(shared_y_test).long()
# if cuda:
# shared_x_test = tuple(d.cuda() for d in shared_x_test)
# shared_y_test = torch.Tensor(shared_y_test).long().cuda()
model1_x_test = torch.stack(model1_x_test)
model2_x_test = torch.stack(model2_x_test)
model1.eval()
shared_x_test = torch.stack(shared_x_test)
model1.eval()
adversary = LinfPGDAttack(
model1, loss_fn=nn.CrossEntropyLoss(reduction="sum"), eps=eps,
nb_iter=adv_steps, eps_iter=0.01, rand_init=True, clip_min=0.0, clip_max=1.0,
targeted=False)
adv_untargeted = adversary.perturb(shared_x_test, shared_y_test)
timestr = time.strftime("%Y%m%d_%H%M%S")
print("saving models at", timestr)
model1_name = './models/{}_{}_{}_model1_{}.pickle'.format(task,num_shared_classes, percent_shared_data,timestr)
model2_name = './models/{}_{}_{}_model2_{}.pickle'.format(task,num_shared_classes, percent_shared_data,timestr)
adv_name = './models/{}_{}_{}_adv_{}.pickle'.format(task,num_shared_classes, percent_shared_data,timestr)
torch.save(model1, model1_name)
torch.save(model2, model2_name)
torch.save(adversary, adv_name)
# Eval
with torch.no_grad():
model1.eval()
model2.eval()
# model1 outputs
output1 = model1(model1_x_test)
shared_output1 = model1(shared_x_test)
adv_output1 = model1(adv_untargeted)
# model2 outputs
output2 = model2(model2_x_test)
shared_output2 = model2(shared_x_test)
adv_output2 = model2(adv_untargeted)
if task.upper() == "CIFAR100":
# model 1
print("model1_acc:", accuracy(output1,model1_y_test))
print("model1_acc_5:", accuracy_n(output1,model1_y_test,5))
print("model1_acc_shared:", accuracy(shared_output1,shared_y_test))
print("model1_acc_5_shared:", accuracy_n(shared_output1,shared_y_test,5))
print("model1_adv_acc_shared:", accuracy(adv_output1,shared_y_test))
print("model1_adv_acc_5_shared:", accuracy_n(adv_output1,shared_y_test,5))
print()
# model 2
print("model2_acc:", accuracy(output2,model2_y_test))
print("model2_acc_5:", accuracy_n(output2,model2_y_test,5))
print("model2_acc_shared:", accuracy(shared_output2,shared_y_test))
print("model2_acc_5_shared:", accuracy_n(shared_output2,shared_y_test,5))
print("model2_adv_acc_shared:", accuracy(adv_output2,shared_y_test))
print("model2_adv_acc_5_shared:", accuracy_n(adv_output2,shared_y_test,5))
else:
# model 1
print("model1_acc:", accuracy(output1,model1_y_test))
print("model1_acc_shared:", accuracy(shared_output1,shared_y_test))
print("model1_adv_acc_shared:", accuracy(adv_output1,shared_y_test))
print()
# model 2
print("model2_acc:", accuracy(output2,model2_y_test))
print("model2_acc_shared:", accuracy(shared_output2,shared_y_test))
print("model2_adv_acc_shared:", accuracy(adv_output2,shared_y_test))
def init_imagenet(dpath=IMAGENET_LOCATION):
from torchvision.datasets import ImageNet
imagenet = ImageNet(IMAGENET_LOCATION, split='val', download=True)
return imagenet
def get_train_val_loaders(
root_path: str,
train_transforms: Callable,
val_transforms: Callable,
batch_size: int = 16,
num_workers: int = 8,
val_batch_size: Optional[int] = None,
limit_train_num_samples: Optional[int] = None,
limit_val_num_samples: Optional[int] = None,
) -> Tuple[DataLoader, DataLoader, DataLoader]:
train_ds = ImageNet(
root_path,
split="train",
transform=lambda sample: train_transforms(image=sample)["image"],
loader=opencv_loader)
val_ds = ImageNet(
root_path,
split="val",
transform=lambda sample: val_transforms(image=sample)["image"],
loader=opencv_loader)
if limit_train_num_samples is not None:
np.random.seed(limit_train_num_samples)
train_indices = np.random.permutation(
len(train_ds))[:limit_train_num_samples]
train_ds = Subset(train_ds, train_indices)
if limit_val_num_samples is not None:
np.random.seed(limit_val_num_samples)
val_indices = np.random.permutation(
len(val_ds))[:limit_val_num_samples]
val_ds = Subset(val_ds, val_indices)
# random samples for evaluation on training dataset
if len(val_ds) < len(train_ds):
np.random.seed(len(val_ds))
train_eval_indices = np.random.permutation(len(train_ds))[:len(val_ds)]
train_eval_ds = Subset(train_ds, train_eval_indices)
else:
train_eval_ds = train_ds
train_loader = idist.auto_dataloader(
train_ds,
shuffle=True,
batch_size=batch_size,
num_workers=num_workers,
drop_last=True,
)
val_batch_size = batch_size * 4 if val_batch_size is None else val_batch_size
val_loader = idist.auto_dataloader(
val_ds,
shuffle=False,
batch_size=val_batch_size,
num_workers=num_workers,
drop_last=False,
)
train_eval_loader = idist.auto_dataloader(
train_eval_ds,
shuffle=False,
batch_size=val_batch_size,
num_workers=num_workers,
drop_last=False,
)
return train_loader, val_loader, train_eval_loader
def main_worker(gpu, ngpus_per_node, args):
args.gpu = gpu
logger = get_logger(args.logging_file)
logger.info("Use GPU: {} for training".format(args.gpu))
args.rank = args.rank * ngpus_per_node + gpu
torch.distributed.init_process_group(backend="nccl",
init_method=args.dist_url,
world_size=args.world_size,
rank=args.rank)
epochs = args.epochs
input_size = args.input_size
resume_epoch = args.resume_epoch
initializer = KaimingInitializer()
zero_gamma = ZeroLastGamma()
mix_precision_training = args.mix_precision_training
is_first_rank = True if args.rank % ngpus_per_node == 0 else False
batches_pre_epoch = args.num_training_samples // (args.batch_size *
ngpus_per_node)
lr = 0.1 * (args.batch_size * ngpus_per_node //
32) if args.lr == 0 else args.lr
model = get_model(models, args.model)
model.apply(initializer)
if args.last_gamma:
model.apply(zero_gamma)
logger.info('Apply zero last gamma init.')
if is_first_rank and args.model_info:
summary(model, torch.rand((1, 3, input_size, input_size)))
parameters = model.parameters() if not args.no_wd else no_decay_bias(model)
if args.sgd_gc:
logger.info('Use SGD_GC optimizer.')
optimizer = SGD_GC(parameters,
lr=lr,
momentum=args.momentum,
weight_decay=args.wd,
nesterov=True)
else:
optimizer = optim.SGD(parameters,
lr=lr,
momentum=args.momentum,
weight_decay=args.wd,
nesterov=True)
lr_scheduler = CosineWarmupLr(optimizer,
batches_pre_epoch,
epochs,
base_lr=args.lr,
warmup_epochs=args.warmup_epochs)
# dropblock_scheduler = DropBlockScheduler(model, batches_pre_epoch, epochs)
if args.lookahead:
optimizer = Lookahead(optimizer)
logger.info('Use lookahead optimizer.')
torch.cuda.set_device(args.gpu)
model.cuda(args.gpu)
args.num_workers = int(
(args.num_workers + ngpus_per_node - 1) / ngpus_per_node)
if args.mix_precision_training and is_first_rank:
logger.info('Train with FP16.')
scaler = GradScaler(enabled=args.mix_precision_training)
model = nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu])
Loss = nn.CrossEntropyLoss().cuda(args.gpu) if not args.label_smoothing else \
LabelSmoothingLoss(args.classes, smoothing=0.1).cuda(args.gpu)
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
if args.autoaugment:
train_transform = transforms.Compose([
transforms.RandomResizedCrop(input_size),
transforms.RandomHorizontalFlip(),
ImageNetPolicy,
transforms.ToTensor(),
normalize,
])
else:
train_transform = transforms.Compose([
transforms.RandomResizedCrop(input_size),
# Cutout(),
transforms.RandomHorizontalFlip(),
transforms.ColorJitter(0.4, 0.4, 0.4),
transforms.ToTensor(),
normalize,
])
val_transform = transforms.Compose([
transforms.Resize(int(input_size / 0.875)),
transforms.CenterCrop(input_size),
transforms.ToTensor(),
normalize,
])
train_set = ImageNet(args.data_path,
split='train',
transform=train_transform)
val_set = ImageNet(args.data_path, split='val', transform=val_transform)
train_sampler = DistributedSampler(train_set)
train_loader = DataLoader(train_set,
args.batch_size,
False,
pin_memory=True,
num_workers=args.num_workers,
drop_last=True,
sampler=train_sampler)
val_loader = DataLoader(val_set,
args.batch_size,
False,
pin_memory=True,
num_workers=args.num_workers,
drop_last=False)
if resume_epoch > 0:
loc = 'cuda:{}'.format(args.gpu)
checkpoint = torch.load(args.resume_param, map_location=loc)
model.load_state_dict(checkpoint['model'])
optimizer.load_state_dict(checkpoint['optimizer'])
scaler.load_state_dict(checkpoint['scaler'])
lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
print("Finish loading resume param.")
torch.backends.cudnn.benchmark = True
top1_acc = metric.Accuracy(name='Top1 Accuracy')
top5_acc = metric.TopKAccuracy(top=5, name='Top5 Accuracy')
loss_record = metric.NumericalCost(name='Loss')
for epoch in range(resume_epoch, epochs):
tic = time.time()
train_sampler.set_epoch(epoch)
if not args.mixup:
train_one_epoch(model, train_loader, Loss, optimizer, epoch,
lr_scheduler, logger, top1_acc, loss_record,
scaler, args)
else:
train_one_epoch_mixup(model, train_loader, Loss, optimizer, epoch,
lr_scheduler, logger, loss_record, scaler,
args)
train_speed = int(args.num_training_samples // (time.time() - tic))
if is_first_rank:
logger.info(
'Finish one epoch speed: {} samples/s'.format(train_speed))
test(model, val_loader, Loss, epoch, logger, top1_acc, top5_acc,
loss_record, args)
if args.rank % ngpus_per_node == 0:
checkpoint = {
'model': model.state_dict(),
'optimizer': optimizer.state_dict(),
'scaler': scaler.state_dict(),
'lr_scheduler': lr_scheduler.state_dict(),
}
torch.save(
checkpoint, '{}/{}_{}_{:.5}.pt'.format(args.save_dir,
args.model, epoch,
top1_acc.get()))
default=100,
help='number of examples/minibatch')
parser.add_argument('--num_batches',
type=int,
required=False,
help='number of batches (default entire dataset)')
args = parser.parse_args()
model = resnet50(pretrained=True, progress=True)
normalizer = utils.DifferentiableNormalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
dataset = ImageNet(
args.imagenet_path,
split='val',
transform=transforms.Compose([
transforms.CenterCrop(224),
transforms.ToTensor(),
]),
)
val_loader = DataLoader(
dataset,
batch_size=args.batch_size,
shuffle=True,
)
model.eval()
if torch.cuda.is_available():
model.cuda()
cw_loss = lf.CWLossF6(model, normalizer, kappa=float('inf'))
perturbation_loss = lf.PerturbationNormLoss(lp=2)
def main(args):
model = Model()
optimizer = torch.optim.SGD(model.parameters(), lr=args.learning_rate,
momentum=args.momentum, weight_decay=args.weight_decay,
nesterov=args.nesterov)
if args.scheduler=='multistep':
scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, args.milestones,
gamma=args.gamma)
elif args.scheduler=='cosine':
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, args.step_ssize)
criterion = torch.nn.CrossEntropyLoss()
model = model.cuda()
criterion = criterion.cuda()
start_epoch = 0
# Check number of parameters your model
pytorch_total_params = sum(p.numel() for p in model.parameters())
print(f"Number of parameters: {pytorch_total_params}")
if not os.path.exists('{}'.format(args.savepath)):
os.makedirs('{}'.format(args.savepath))
# resume
if args.resume:
model, optimizer, start_epoch = load_ckpt(model, optimizer, args)
# Dataloader
if args.dataset=='cifar10':
normalize = transforms.Normalize(
mean=[0.4914, 0.4822, 0.4465],
std=[0.2023, 0.1994, 0.2010])
transform_train = transforms.Compose([
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize,
])
transform_val = transforms.Compose([
transforms.ToTensor(),
normalize,
])
trainset = CIFAR10(
root=args.datapath, train=True, download=True,
transform=transform_train)
valset = CIFAR10(
root=args.datapath, train=False, download=True,
transform=transform_val)
elif args.dataset=='cifar100':
normalize = transforms.Normalize(
mean=[0.4914, 0.4822, 0.4465],
std=[0.2023, 0.1994, 0.2010])
transform_train = transforms.Compose([
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize,
])
transform_val = transforms.Compose([
transforms.ToTensor(),
normalize,
])
trainset = CIFAR100(
root=args.datapath, train=True, download=True,
transform=transform_train)
valset = CIFAR100(
root=args.datapath, train=False, download=True,
transform=transform_val)
elif args.dataset=='ImageNet':
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
transform_train = transforms.Compose([
transforms.RandomResizedCrop(image_size),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize,
])
transform_val = transforms.Compose([
transforms.Resize(image_size + 32),
transforms.CenterCrop(image_size),
transforms.ToTensor(),
normalize,
])
trainset = ImageNet(
root=args.datapath, split='train', download=False,
transform=transform_train)
valset = ImageNet(
root=args.datapath, split='val', download=False,
transform=transform_val)
elif args.dataeset=='tiny-imagenet-200':
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
transform_train = transforms.Compose([
transforms.RandomResizedCrop(image_size),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize,
])
transform_val = transforms.Compose([
transforms.Resize(image_size + 32),
transforms.CenterCrop(image_size),
transforms.ToTensor(),
normalize,
])
trainset = ImageFolder(
root=args.datapath, split='train', download=False,
transform=transform_train)
valset = ImageFolder(
root=args.datapath, split='val', download=False,
transform=transform_val)
train_loader = torch.utils.data.DataLoader(
trainset,
batch_size=args.batch_size, shuffle=True,
num_workers=args.num_workers, pin_memory=False)
val_loader = torch.utils.data.DataLoader(
valset,
batch_size=args.batch_size, shuffle=False,
num_workers=args.num_workers, pin_memory=False)
# start training
last_top1_acc = 0
acc1_valid = 0
best_acc1 = 0
is_best = False
for epoch in range(start_epoch, args.epochs):
print("\n----- epoch: {}, lr: {} -----".format(
epoch, optimizer.param_groups[0]["lr"]))
# train for one epoch
start_time = time.time()
last_top1_acc = train(train_loader, epoch, model, optimizer, criterion)
elapsed_time = time.time() - start_time
print('==> {:.2f} seconds to train this epoch\n'.format(
elapsed_time))
# validate for one epoch
start_time = time.time()
acc1_valid = validate(val_loader, model, criterion)
elapsed_time = time.time() - start_time
print('==> {:.2f} seconds to validate this epoch\n'.format(
elapsed_time))
# learning rate scheduling
scheduler.step()
is_best = acc1_valid > best_acc1
best_acc1 = max(acc1_valid, best_acc1)
checkpoint = {
'epoch': epoch + 1,
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict()
}
save_ckpt(checkpoint, is_best, args)
#if is_best:
# torch.save(model.state_dict(), args.savepath+'model_weight_best.pth')
# Save model each epoch
#torch.save(model.state_dict(), args.savepath+'model_weight_epoch{}.pth'.format(epoch))
print(f"Last Top-1 Accuracy: {last_top1_acc}")
print(f"Number of parameters: {pytorch_total_params}")
def __len__(self):
return ImageNet.__len__(self)
else:
optimizer = torch.optim.SGD(get_parameters(model),
lr=0.05,
momentum=momentum,
nesterov=nesterov,
weight_decay=4e-5)
if args.train_mode == 'search':
from nni.retiarii.oneshot.pytorch import ProxylessTrainer
from torchvision.datasets import ImageNet
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
dataset = ImageNet(args.data_path,
transform=transforms.Compose([
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize,
]))
trainer = ProxylessTrainer(model,
loss=LabelSmoothingLoss(),
dataset=dataset,
optimizer=optimizer,
metrics=lambda output, target: accuracy(
output, target, topk=(
1,
5,
)),
num_epochs=120,
log_frequency=10)
trainer.fit()
def create_dataset(name,
root,
split='validation',
search_split=True,
class_map=None,
load_bytes=False,
is_training=False,
download=False,
batch_size=None,
repeats=0,
**kwargs):
""" Dataset factory method
In parenthesis after each arg are the type of dataset supported for each arg, one of:
* folder - default, timm folder (or tar) based ImageDataset
* torch - torchvision based datasets
* TFDS - Tensorflow-datasets wrapper in IterabeDataset interface via IterableImageDataset
* all - any of the above
Args:
name: dataset name, empty is okay for folder based datasets
root: root folder of dataset (all)
split: dataset split (all)
search_split: search for split specific child fold from root so one can specify
`imagenet/` instead of `/imagenet/val`, etc on cmd line / config. (folder, torch/folder)
class_map: specify class -> index mapping via text file or dict (folder)
load_bytes: load data, return images as undecoded bytes (folder)
download: download dataset if not present and supported (TFDS, torch)
is_training: create dataset in train mode, this is different from the split.
For Iterable / TDFS it enables shuffle, ignored for other datasets. (TFDS)
batch_size: batch size hint for (TFDS)
repeats: dataset repeats per iteration i.e. epoch (TFDS)
**kwargs: other args to pass to dataset
Returns:
Dataset object
"""
name = name.lower()
if name.startswith('torch/'):
name = name.split('/', 2)[-1]
torch_kwargs = dict(root=root, download=download, **kwargs)
if name in _TORCH_BASIC_DS:
ds_class = _TORCH_BASIC_DS[name]
use_train = split in _TRAIN_SYNONYM
ds = ds_class(train=use_train, **torch_kwargs)
elif name == 'inaturalist' or name == 'inat':
assert has_inaturalist, 'Please update to PyTorch 1.10, torchvision 0.11+ for Inaturalist'
target_type = 'full'
split_split = split.split('/')
if len(split_split) > 1:
target_type = split_split[0].split('_')
if len(target_type) == 1:
target_type = target_type[0]
split = split_split[-1]
if split in _TRAIN_SYNONYM:
split = '2021_train'
elif split in _EVAL_SYNONYM:
split = '2021_valid'
ds = INaturalist(version=split,
target_type=target_type,
**torch_kwargs)
elif name == 'places365':
assert has_places365, 'Please update to a newer PyTorch and torchvision for Places365 dataset.'
if split in _TRAIN_SYNONYM:
split = 'train-standard'
elif split in _EVAL_SYNONYM:
split = 'val'
ds = Places365(split=split, **torch_kwargs)
elif name == 'imagenet':
if split in _EVAL_SYNONYM:
split = 'val'
ds = ImageNet(split=split, **torch_kwargs)
elif name == 'image_folder' or name == 'folder':
# in case torchvision ImageFolder is preferred over timm ImageDataset for some reason
if search_split and os.path.isdir(root):
# look for split specific sub-folder in root
root = _search_split(root, split)
ds = ImageFolder(root, **kwargs)
else:
assert False, f"Unknown torchvision dataset {name}"
elif name.startswith('tfds/'):
ds = IterableImageDataset(root,
parser=name,
split=split,
is_training=is_training,
download=download,
batch_size=batch_size,
repeats=repeats,
**kwargs)
else:
# FIXME support more advance split cfg for ImageFolder/Tar datasets in the future
if search_split and os.path.isdir(root):
# look for split specific sub-folder in root
root = _search_split(root, split)
ds = ImageDataset(root,
parser=name,
class_map=class_map,
load_bytes=load_bytes,
**kwargs)
return ds
def get_train_val_loaders(
root_path: str,
train_transforms: Callable,
val_transforms: Callable,
batch_size: int = 16,
num_workers: int = 8,
val_batch_size: Optional[int] = None,
pin_memory: bool = True,
random_seed: Optional[int] = None,
train_sampler: Optional[Union[Sampler, str]] = None,
val_sampler: Optional[Union[Sampler, str]] = None,
limit_train_num_samples: Optional[int] = None,
limit_val_num_samples: Optional[int] = None,
) -> Tuple[DataLoader, DataLoader, DataLoader]:
train_ds = ImageNet(
root_path,
split="train",
transform=lambda sample: train_transforms(image=sample)["image"],
loader=opencv_loader)
val_ds = ImageNet(
root_path,
split="val",
transform=lambda sample: val_transforms(image=sample)["image"],
loader=opencv_loader)
if limit_train_num_samples is not None:
if random_seed is not None:
np.random.seed(random_seed)
train_indices = np.random.permutation(
len(train_ds))[:limit_train_num_samples]
train_ds = Subset(train_ds, train_indices)
if limit_val_num_samples is not None:
val_indices = np.random.permutation(
len(val_ds))[:limit_val_num_samples]
val_ds = Subset(val_ds, val_indices)
# random samples for evaluation on training dataset
if len(val_ds) < len(train_ds):
train_eval_indices = np.random.permutation(len(train_ds))[:len(val_ds)]
train_eval_ds = Subset(train_ds, train_eval_indices)
else:
train_eval_ds = train_ds
if isinstance(train_sampler, str):
assert train_sampler == "distributed"
train_sampler = data_dist.DistributedSampler(train_ds)
train_eval_sampler = None
if isinstance(val_sampler, str):
assert val_sampler == "distributed"
val_sampler = data_dist.DistributedSampler(val_ds, shuffle=False)
train_eval_sampler = data_dist.DistributedSampler(train_eval_ds,
shuffle=False)
train_loader = DataLoader(
train_ds,
shuffle=train_sampler is None,
batch_size=batch_size,
num_workers=num_workers,
sampler=train_sampler,
pin_memory=pin_memory,
drop_last=True,
)
val_batch_size = batch_size * 4 if val_batch_size is None else val_batch_size
val_loader = DataLoader(
val_ds,
shuffle=False,
sampler=val_sampler,
batch_size=val_batch_size,
num_workers=num_workers,
pin_memory=pin_memory,
drop_last=False,
)
train_eval_loader = DataLoader(
train_eval_ds,
shuffle=False,
sampler=train_eval_sampler,
batch_size=val_batch_size,
num_workers=num_workers,
pin_memory=pin_memory,
drop_last=False,
)
return train_loader, val_loader, train_eval_loader
import torch
from torch.utils.data import DataLoader
from torchvision.datasets import ImageNet
print("Pytorch version:{}".format(torch.__version__))
#useful links
#https://pytorch.org/docs/stable/_modules/torchvision/datasets/imagenet.html#ImageNet
#https://pytorch.org/docs/stable/torchvision/datasets.html#imagenet
#%%
"""
#_________________________________Dataloader for ImageNet______________________________#
# #
# The source code of the class ImageNet from pytorch says #
# #
# The dataset is no longer publicly accessible. #
# You need to download the archives externally and place them in the root directory. #
# #
# I couldn't download the image because it takes 5 days to check a new account on #
# the ImageNet website so i just implement the dataloader with the class ImageNet #
# from pytorch #
# #
#______________________________________________________________________________________#
"""
root_dir = "path_to_ImageNet_directory_file"
data = ImageNet(root=root_dir)
data_loader = DataLoader(data, batch_size=1)
# Model
model = vit_base_patch16_224(pretrained=True).cuda()
baselines = Baselines(model)
# LRP
model_LRP = vit_LRP(pretrained=True).cuda()
model_LRP.eval()
lrp = LRP(model_LRP)
# orig LRP
model_orig_LRP = vit_orig_LRP(pretrained=True).cuda()
model_orig_LRP.eval()
orig_lrp = LRP(model_orig_LRP)
# Dataset loader for sample images
transform = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
])
imagenet_ds = ImageNet(args.imagenet_validation_path, split='val', download=False, transform=transform)
sample_loader = torch.utils.data.DataLoader(
imagenet_ds,
batch_size=args.batch_size,
shuffle=False,
num_workers=4
)
compute_saliency_and_save(args)
def _make_data_loader_imagenet(root,
batch_size,
workers=4,
is_train=True,
download=False,
distributed=False):
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
logger = logging.getLogger('octconv')
if is_train:
logger.info("Loading ImageNet training data")
st = time.time()
scale = (0.08, 1.0)
transform = transforms.Compose([
transforms.RandomResizedCrop(224, scale=scale),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize,
])
dataset = ImageNet(root=root,
split='train',
download=download,
transform=transform)
logger.info("Took: {}".format(time.time() - st))
if distributed:
sampler = torch.utils.data.distributed.DistributedSampler(dataset)
else:
sampler = torch.utils.data.RandomSampler(dataset)
loader = DataLoader(dataset,
batch_size=batch_size,
num_workers=workers,
sampler=sampler,
pin_memory=True)
else:
logger.info("Loading ImageNet validation data")
transform = transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
normalize,
])
dataset = ImageNet(root=root,
split='val',
download=download,
transform=transform)
if distributed:
sampler = torch.utils.data.distributed.DistributedSampler(dataset)
else:
sampler = torch.utils.data.SequentialSampler(dataset)
loader = DataLoader(dataset,
batch_size=batch_size,
num_workers=workers,
sampler=sampler,
pin_memory=True)
return loader
args.batch_size = 1
model = vgg19(pretrained=True).to(device)
args.BP = VanillaBackprop(model)
elif args.method == 'integrad':
args.batch_size = 1
model = vgg19(pretrained=True).to(device)
args.BP = IntegratedGradients(model)
elif args.method == 'gradcam':
args.batch_size = 1
model = vgg19(pretrained=True).to(device)
args.BP = GradCam(model, target_layer=36) # True for VGG-19
else:
raise Exception('No model found')
# Dataset loader for sample images
transform = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor()
])
imagenet_ds = ImageNet('path/to/imagenet/', split='val', download=False,
transform=transform)
sample_loader = torch.utils.data.DataLoader(
imagenet_ds,
batch_size=args.batch_size,
shuffle=False,
num_workers=4
)
generate_visualization_to_hdf5(args)
def get_dataset(path: Path, transforms: Callable = None) -> Dataset:
return ImageNet(str(path), split='val', transform=transforms)
def get_dataset(dataset='mnist',
dataroot='data/',
imresize=None,
augment=False,
normalize=False,
tfm=None):
num_chs = 1 if dataset.lower() in ['rsna', 'mnist'] else 3
valid_dataset = None
if tfm is None:
transform_lst = []
if imresize is not None:
transform_lst.append(transforms.Resize(imresize))
transform_lst.append(transforms.ToTensor())
if normalize:
transform_lst.append(
transforms.Normalize(num_chs * (0.5, ), num_chs * (0.5, )))
transform = transforms.Compose(transform_lst)
else:
transform = tfm
if dataset.lower() == 'rsna':
raise NotImplementedError
elif dataset.lower() == 'imagenet':
train_dataset = ImageNet(dataroot,
train=True,
transform=transform,
download=True)
test_dataset = ImageNet(dataroot,
train=False,
transform=transform,
download=True)
elif dataset.lower() == 'cifar10':
if tfm is None:
transform = transform
normalize_cifar10 = transforms.Normalize(
mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
if normalize: # overwrites the normalize
transform = transforms.Compose([
transforms.ToTensor(),
normalize_cifar10,
])
# else:
# transform_train = transform
train_dataset = CIFAR10(dataroot,
train=True,
transform=transform,
download=True)
test_dataset = CIFAR10(dataroot,
train=False,
transform=transform,
download=True)
elif dataset.lower() == 'cifar100':
train_dataset = CIFAR100(dataroot,
train=True,
transform=transform,
download=True)
test_dataset = CIFAR100(dataroot,
train=False,
transform=transform,
download=True)
elif dataset.lower() == 'mnist':
train_dataset = MNIST(dataroot,
train=True,
transform=transform,
download=True)
test_dataset = MNIST(dataroot,
train=False,
transform=transform,
download=True)
elif dataset.lower() == 'svhn':
dataroot = os.path.join(dataroot, 'svhn')
train_dataset = SVHN(dataroot,
split='train',
transform=transform,
download=True)
test_dataset = SVHN(dataroot,
split='test',
transform=transform,
download=True)
return train_dataset, test_dataset, num_chs
