def get_dataset(args,
config,
test=False,
rev=False,
one_hot=True,
subset=False,
shuffle=True):
total_labels = 10 if config.data.dataset.lower().split(
'_')[0] != 'cifar100' else 100
reduce_labels = total_labels != config.n_labels
if config.data.dataset.lower() in [
'mnist_transferbaseline', 'cifar10_transferbaseline',
'fashionmnist_transferbaseline', 'cifar100_transferbaseline'
]:
print('loading baseline transfer dataset')
rev = True
test = False
subset = True
reduce_labels = True
if config.data.random_flip is False:
transform = transforms.Compose(
[transforms.Resize(config.data.image_size),
transforms.ToTensor()])
else:
if not test:
transform = transforms.Compose([
transforms.Resize(config.data.image_size),
transforms.RandomHorizontalFlip(p=0.5),
transforms.ToTensor()
])
else:
transform = transforms.Compose([
transforms.Resize(config.data.image_size),
transforms.ToTensor()
])
if config.data.dataset.lower().split('_')[0] == 'mnist':
dataset = MNIST(os.path.join(args.run, 'datasets'),
train=not test,
download=True,
transform=transform)
elif config.data.dataset.lower().split('_')[0] in [
'fashionmnist', 'fmnist'
]:
dataset = FashionMNIST(os.path.join(args.run, 'datasets'),
train=not test,
download=True,
transform=transform)
elif config.data.dataset.lower().split('_')[0] == 'cifar10':
dataset = CIFAR10(os.path.join(args.run, 'datasets'),
train=not test,
download=True,
transform=transform)
elif config.data.dataset.lower().split('_')[0] == 'cifar100':
dataset = CIFAR100(os.path.join(args.run, 'datasets'),
train=not test,
download=True,
transform=transform)
else:
raise ValueError('Unknown config dataset {}'.format(
config.data.dataset))
if type(dataset.targets) is list:
# CIFAR10 and CIFAR100 store targets as list, unlike (F)MNIST which uses torch.Tensor
dataset.targets = np.array(dataset.targets)
if not rev:
labels_to_consider = np.arange(config.n_labels)
target_transform = lambda label: single_one_hot_encode(
label, n_labels=config.n_labels)
cond_size = config.n_labels
else:
labels_to_consider = np.arange(config.n_labels, total_labels)
target_transform = lambda label: single_one_hot_encode_rev(
label, start_label=config.n_labels, n_labels=total_labels)
cond_size = total_labels - config.n_labels
if reduce_labels:
idx = np.any(
[np.array(dataset.targets) == i for i in labels_to_consider],
axis=0).nonzero()
dataset.targets = dataset.targets[idx]
dataset.data = dataset.data[idx]
if one_hot:
dataset.target_transform = target_transform
if subset and args.subset_size != 0:
dataset = torch.utils.data.Subset(dataset, np.arange(args.subset_size))
dataloader = DataLoader(dataset,
batch_size=config.training.batch_size,
shuffle=shuffle,
num_workers=0)
return dataloader, dataset, cond_size
def main():
"""Do stuff."""
args = FLAGS.parse_args()
batch_size = 32
sample_size = 100
training_epochs = 20
if args.mode == 'train':
ckpt = torch.load(args.loadname, map_location=device)
model = ckpt['model']
previous_masks = ckpt['previous_masks']
model.add_dataset(args.dataset, args.dataset.num_outputs)
model.set_dataset(args.dataset)
current_dataset_idx = len(model.datasets)
unrelated_tasks = None
previous_samples = None
if current_dataset_idx > 2:
unrelated_tasks = ckpt['unrelated_tasks']
previous_samples = ckpt['previous_samples']
if current_dataset_idx == 3:
previous_samples = ckpt['previous_samples'].cpu()
manager = Manager(model, 0.75, previous_masks, current_dataset_idx,
previous_samples, unrelated_tasks)
if current_dataset_idx == 2:
manager.pruner.current_masks = previous_masks
else:
manager.pruner.initialize_new_mask()
if 'cropped' in args.train_path:
train_loader = train_loader_cropped(args.train_path, batch_size)
test_loader = test_loader_cropped(args.test_path, batch_size)
sample_loader = train_loader_sample(args.train_path, sample_size)
if args.dataset == 'mnist':
mnist_trainset = MNIST(
root='./data',
train=True,
download=True,
transform=transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Lambda(lambda x: x.repeat(3, 1, 1)),
transforms.Normalize((0.1307, ), (0.3081, ))
]))
mnist_testset = MNIST(
root='./data',
train=False,
download=True,
transform=transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Lambda(lambda x: x.repeat(3, 1, 1)),
transforms.Normalize((0.1307, ), (0.3081, ))
]))
mnist_sampleset = MNIST(root='./data',
train=True,
download=True,
transform=transform_sample())
train_loader = data.DataLoader(mnist_trainset,
batch_size,
shuffle=True)
test_loader = data.DataLoader(mnist_testset, batch_size)
sample_loader = data.DataLoader(mnist_sampleset,
sample_size,
shuffle=True)
if args.dataset == 'fashion_mnist':
fmnist_trainset = FashionMNIST(
root='./data',
train=True,
download=True,
transform=transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(), # image to Tensor
transforms.Lambda(lambda x: x.repeat(3, 1, 1)),
transforms.Normalize((0.1307, ),
(0.3081, )) # image, label
]))
fmnist_testset = FashionMNIST(
root='./data',
train=False,
download=True,
transform=transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Lambda(lambda x: x.repeat(3, 1, 1)),
transforms.Normalize((0.1307, ), (0.3081, ))
]))
fmnist_sampleset = FashionMNIST(root='./data',
train=True,
download=True,
transform=transform_sample())
train_loader = data.DataLoader(fmnist_trainset,
batch_size,
shuffle=True)
test_loader = data.DataLoader(fmnist_testset, batch_size)
sample_loader = data.DataLoader(fmnist_sampleset,
sample_size,
shuffle=True)
else:
train_loader = train_loader_noncropped(args.train_path, batch_size)
test_loader = test_loader_noncropped(args.test_path, batch_size)
sample_loader = train_loader_sample(args.train_path, sample_size)
manager.train(dataset_idx=current_dataset_idx,
epochs=training_epochs,
train_loader=train_loader,
sample_loader=sample_loader,
isRetrain=False)
manager.prune(current_dataset_idx, train_loader)
manager.eval_task_specific(current_dataset_idx, test_loader)
if args.mode == 'eval':
ckpt = torch.load(args.loadname, map_location=device)
model = ckpt['model']
previous_masks = ckpt['previous_masks']
unrelated_tasks = None
previous_samples = None
current_dataset_idx = model.datasets.index(args.dataset) + 1
if len(model.datasets) > 2:
unrelated_tasks = ckpt['unrelated_tasks']
previous_samples = ckpt['previous_samples']
model.set_dataset(args.dataset)
if 'cropped' in args.test_path:
test_loader = test_loader_cropped(args.test_path, batch_size)
if args.dataset == 'mnist':
mnist_testset = MNIST(
root='./data',
train=False,
download=True,
transform=transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Lambda(lambda x: x.repeat(3, 1, 1)),
transforms.Normalize((0.1307, ), (0.3081, ))
]))
test_loader = data.DataLoader(mnist_testset, batch_size)
if args.dataset == 'fashion_mnist':
fmnist_testset = FashionMNIST(
root='./data',
train=False,
download=True,
transform=transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Lambda(lambda x: x.repeat(3, 1, 1)),
transforms.Normalize((0.1307, ), (0.3081, ))
]))
test_loader = data.DataLoader(fmnist_testset, batch_size)
else:
test_loader = test_loader_noncropped(args.test_path, batch_size)
manager = Manager(model, 0.75, previous_masks, current_dataset_idx,
previous_samples, unrelated_tasks)
manager.eval_task_specific(current_dataset_idx, test_loader)
def experiment(num_shared_classes,
percent_shared_data,
n_epochs=200,
batch_size=128,
eps=.3,
adv_steps=1000,
learning_rate=.0004,
gpu_num=1,
adv_training="none",
task="CIFAR100",
masked=False,
savemodel=False,
download_data=False):
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()
if task.upper() == "CIFAR100":
mean = (0.5070751592371323, 0.48654887331495095, 0.4409178433670343)
std = (0.2673342858792401, 0.2564384629170883, 0.27615047132568404)
transform_test = transforms.Compose(
[transforms.ToTensor(),
transforms.Normalize(mean, std)])
transform_train = transforms.Compose([
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize(mean, std)
])
train_data = CIFAR100("data/",
transform=transform_train,
download=download_data)
test_data = CIFAR100("data/",
train=False,
transform=transform_test,
download=download_data)
elif task.upper() == "FASHIONMNIST":
mean = (0.2860)
std = (0.3530)
transform = transforms.Compose([
transforms.Lambda(lambda image: image.convert('RGB')),
transforms.ToTensor()
])
train_data = FashionMNIST('data/fashionmnist',
transform=transform,
train=True,
download=download_data)
test_data = FashionMNIST('data/fashionmnist',
transform=transform,
train=False,
download=download_data)
elif task.upper() == "IMAGENET":
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
transform_train = transforms.Compose([
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize,
])
transform_test = transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
normalize,
])
train_data = ImageFolder(args.data + '/train',
transform=transform_train)
#################################### change to ImageFolder instead of CustomImageFolderTest
test_data = ImageFolder(args.data + '/val/', transform=transform_test)
else:
mean = (0.4914, 0.4822, 0.4465)
std = (0.2470, 0.2435, 0.2616)
transform_test = transforms.Compose(
[transforms.ToTensor(),
transforms.Normalize(mean, std)])
transform_train = transforms.Compose([
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize(mean, std),
])
train_data = CIFAR10("data/",
transform=transform_train,
download=download_data)
test_data = CIFAR10("data/",
train=False,
transform=transform_test,
download=download_data)
######################### Change all_classes so it's not looping over 1.2 million images
if task.upper() == "IMAGENET":
all_classes = range(len(train_data.classes))
else:
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))
# split
############################## Change the way classes_by_index is generated to avoid looping over dataset
if task.upper() == "IMAGENET":
classes_by_index = []
for idx, label in enumerate(train_data.classes):
label_size = len([
x for x in os.listdir(os.path.join(args.data, "train", label))
if ".JPEG" in x
])
for i in range(label_size):
classes_by_index.append(idx)
classes_by_index = np.array(classes_by_index)
else:
classes_by_index = np.array(
[train_data[index][1] for index in range(len(train_data))])
model1_train_indicies = np.array([])
model2_train_indicies = np.array([])
# cut back on percentage of data
if percent_shared_data < 100:
d = defaultdict(list)
for i in range(len(classes_by_index)):
d[classes_by_index[i]].append(i)
model1_train_indicies = np.array([])
model2_train_indicies = np.array([])
for key in d.keys():
if key in model1_classes:
np.random.shuffle(d[key])
new_len = len(d[key]) // 2
model1_train_indicies = np.concatenate(
(d[key][:new_len], model1_train_indicies))
if key in model2_classes:
np.random.shuffle(d[key])
new_len = len(d[key]) // 2
model2_train_indicies = np.concatenate(
(d[key][:new_len], model2_train_indicies))
else:
model1_train_indicies = np.argwhere(
np.isin(classes_by_index, model1_classes) == True)
model2_train_indicies = np.argwhere(
np.isin(classes_by_index, model2_classes) == True)
# split up shared data
# divide shared data
if percent_shared_data < 100:
# split based on shared classes
for curr_class in shared_classes:
temp_x = np.argwhere(classes_by_index == curr_class).reshape(-1)
size_shared = int(len(temp_x) * (percent_shared_data / 100))
np.random.shuffle(temp_x)
# joint shared data
shared_indicies = temp_x[:size_shared]
# add the joint data to datasets
model1_train_indicies = np.concatenate(
(model1_train_indicies.reshape(-1), shared_indicies))
model2_train_indicies = np.concatenate(
(model2_train_indicies.reshape(-1), shared_indicies))
# disjoint shared class data
point = (len(temp_x) - size_shared) // 2
model1_train_indicies = np.concatenate(
(model1_train_indicies.reshape(-1),
temp_x[size_shared:size_shared + point]))
model2_train_indicies = np.concatenate(
(model2_train_indicies.reshape(-1),
temp_x[size_shared + point:]))
else:
shared_data_indicies = np.argwhere(
np.isin(classes_by_index, shared_classes) == True).reshape(-1)
model1_train_indicies = np.concatenate(
(model1_train_indicies.reshape(-1), shared_data_indicies))
model2_train_indicies = np.concatenate(
(model2_train_indicies.reshape(-1), shared_data_indicies))
# create class mappings
model1_class_mapping = {}
model2_class_mapping = {}
model1_classes_inc = 0
# go through model1 and assign unique classes to incrimental int starting at 0
for c in (shared_classes + model1_classes):
# if it doesn't exist assign
if c not in model1_class_mapping.keys():
model1_class_mapping[c] = model1_classes_inc
model1_classes_inc += 1
model2_classes_inc = 0
# go through model2 and assign unique classes to incrimental int starting at 0
for c in (shared_classes + model2_classes):
# if it doesn't exist in model2 OR in model1, assign it
if c not in model2_class_mapping.keys(
) and c not in model1_class_mapping.keys():
model2_class_mapping[c] = model2_classes_inc
model2_classes_inc += 1
if c in model1_class_mapping.keys():
model2_class_mapping[c] = model1_class_mapping[c]
model1_classes_len = len(model1_class_mapping.keys())
model2_classes_len = len(model2_class_mapping.keys())
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":
if args.model == 'resnet18':
model1 = models.resnet18()
model2 = models.resnet18()
model1.fc = nn.Linear(512, model1_classes_len)
model2.fc = nn.Linear(512, model2_classes_len)
elif args.model == 'resnet50':
model1 = models.resnet50()
model2 = models.resnet50()
model1.fc = nn.Linear(2048, model1_classes_len)
model2.fc = nn.Linear(2048, model2_classes_len)
elif args.model == 'resnet50_2':
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)
################ Changed way cuda device was called and add DataParallel for the models
device = 'cuda' if torch.cuda.is_available() else 'cpu'
if device == 'cuda':
model1 = torch.nn.DataParallel(model1,
device_ids=range(
torch.cuda.device_count()))
model2 = torch.nn.DataParallel(model2,
device_ids=range(
torch.cuda.device_count()))
# Model Training
model1 = model1.to(device)
model2 = model2.to(device)
if task.upper() != "IMAGENET":
criterion1 = nn.CrossEntropyLoss().cuda()
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().cuda()
optimizer2 = optim.AdamW(model2.parameters(), lr=learning_rate)
scheduler2 = optim.lr_scheduler.MultiStepLR(
optimizer2, milestones=[60, 120,
160], gamma=.2) #learning rate decay
else:
onecycle1 = OneCycle.OneCycle(int(
len(model1_train_indicies) * n_epochs / batch_size),
0.8,
prcnt=(n_epochs - 82) * 100 / n_epochs,
momentum_vals=(0.95, 0.8))
onecycle2 = OneCycle.OneCycle(int(
len(model2_train_indicies) * n_epochs / batch_size),
0.8,
prcnt=(n_epochs - 82) * 100 / n_epochs,
momentum_vals=(0.95, 0.8))
criterion1 = nn.CrossEntropyLoss()
optimizer1 = optim.SGD(model1.parameters(),
lr=learning_rate,
momentum=0.95,
weight_decay=1e-4)
criterion2 = nn.CrossEntropyLoss()
optimizer2 = optim.SGD(model2.parameters(),
lr=learning_rate,
momentum=0.95,
weight_decay=1e-4)
# make datasets
model1_train_dataset = SplitDataset(model1_train_indicies, train_data,
model1_class_mapping)
model2_train_dataset = SplitDataset(model2_train_indicies, train_data,
model2_class_mapping)
# create trainloader 1
trainloader_1 = torch.utils.data.DataLoader(model1_train_dataset,
batch_size=batch_size,
shuffle=True,
pin_memory=True,
num_workers=args.workers)
# create trainloader 2
trainloader_2 = torch.utils.data.DataLoader(model2_train_dataset,
batch_size=batch_size,
shuffle=True,
pin_memory=True,
num_workers=args.workers)
# get test sets ready
############################## Change the way test_classes_by_index is generated to avoid looping over dataset
if task.upper() == "IMAGENET":
test_classes_by_index = []
for idx, label in enumerate(test_data.classes):
label_size = len([
x for x in os.listdir(os.path.join(args.data, "val", label))
if ".JPEG" in x
])
for i in range(label_size):
test_classes_by_index.append(idx)
test_classes_by_index = np.array(test_classes_by_index)
else:
test_classes_by_index = np.array(
[test_data[index][1] for index in range(len(test_data))])
model1_test_indicies = np.array([])
model2_test_indicies = np.array([])
shared_test_indicies = np.array([])
model1_test_indicies = np.argwhere(
np.isin(test_classes_by_index, model1_classes) == True)
model2_test_indicies = np.argwhere(
np.isin(test_classes_by_index, model2_classes) == True)
shared_test_indicies = np.argwhere(
np.isin(test_classes_by_index, shared_classes) == True)
model1_test_indicies = np.concatenate(
[model1_test_indicies, shared_test_indicies])
model2_test_indicies = np.concatenate(
[model2_test_indicies, shared_test_indicies])
model1_test_dataset = SplitDataset(model1_test_indicies, test_data,
model1_class_mapping)
model2_test_dataset = SplitDataset(model2_test_indicies, test_data,
model2_class_mapping)
shared_test_dataset = SplitDataset(
shared_test_indicies, test_data,
model2_class_mapping) # does not matter which mapping
# # dataloaders
testloader_1 = torch.utils.data.DataLoader(
model1_test_dataset,
batch_size=len(model1_test_dataset),
shuffle=True,
pin_memory=True,
num_workers=args.workers)
testloader_2 = torch.utils.data.DataLoader(
model2_test_dataset,
batch_size=len(model2_test_dataset),
shuffle=True,
pin_memory=True,
num_workers=args.workers)
testloader_shared = torch.utils.data.DataLoader(shared_test_dataset,
batch_size=batch_size,
shuffle=True,
pin_memory=True,
num_workers=args.workers)
# shared_x_test = []
# shared_y_test = []
# for i in range(len(shared_test_dataset)):
# data = shared_test_dataset[i]
# shared_x_test.append(data[0])
# shared_y_test.append(data[1])
#
# shared_x_test = torch.stack(shared_x_test)
# shared_y_test = torch.tensor(shared_y_test)
if adv_training == "fast_free":
model1 = fast_free_adv_training(model1,
trainloader_1,
epochs=n_epochs,
mean=mean,
std=std,
device=device,
early_stop=False)
print('Finished Training model1 adversarially')
model2 = fast_free_adv_training(model2,
trainloader_2,
epochs=n_epochs,
mean=mean,
std=std,
device=device,
early_stop=False)
print('Finished Training model2 adversarially')
# Adversarial Training for Free! (adapted from https://github.com/mahyarnajibi/FreeAdversarialTraining)
elif adv_training == "free":
# mean = torch.tensor(mean).expand(3,32,32).to(device)
# std = torch.tensor(std).expand(3,32,32).to(device)
global_noise_data = torch.zeros([batch_size, 3, 32, 32]).to(device)
# train model 1
for epoch in tqdm(
range(n_epochs),
desc="Epoch"): # loop over the dataset multiple times
for i, data in enumerate(trainloader_1, 0):
#if cuda: ################################################################################check here if broke
#data = tuple(d.cuda() for d in data)
# get the inputs; data is a list of [inputs, labels]
inputs, labels = data
# steps for adv training
for j in range(4):
noise_batch = Variable(global_noise_data[0:inputs.size(0)],
requires_grad=True).to(device)
in1 = inputs + noise_batch
in1.clamp_(0, 1.0)
in1 = normalize_imgs(in1, mean, std)
output = model1(in1)
loss = criterion1(output, labels)
# compute gradient and do SGD step
optimizer1.zero_grad()
loss.backward()
# Update the noise for the next iteration
pert = fgsm(noise_batch.grad, 4)
global_noise_data[0:inputs.size(0)] += pert.data
global_noise_data.clamp_(-4.0, 4.0)
optimizer1.step()
global_noise_data = torch.zeros([batch_size, 3, 32, 32]).to(device)
# train model 2
for epoch in tqdm(
range(n_epochs),
desc="Epoch"): # loop over the dataset multiple times
for i, data in enumerate(trainloader_2, 0):
#if cuda: ################################################################################check here if broke
# data = tuple(d.cuda() for d in data)
# get the inputs; data is a list of [inputs, labels]
inputs, labels = data
# steps for adv training
for j in range(4):
noise_batch = Variable(global_noise_data[0:inputs.size(0)],
requires_grad=True).to(device)
in1 = inputs + noise_batch
in1.clamp_(0, 1.0)
in1 = normalize_imgs(in1, mean, std)
output = model2(in1)
loss = criterion2(output, labels)
# compute gradient and do SGD step
optimizer2.zero_grad()
loss.backward()
# Update the noise for the next iteration
pert = fgsm(noise_batch.grad, 4)
global_noise_data[0:inputs.size(0)] += pert.data
global_noise_data.clamp_(-4.0, 4.0)
optimizer2.step()
else:
# 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: ################################################################################check here if broke
# data = tuple(d.cuda() for d in data)
# get the inputs; data is a list of [inputs, labels]
inputs, labels = data
# one cycle policy
if task.upper() == "IMAGENET":
lr, mom = onecycle1.calc()
update_lr(optimizer1, lr)
update_mom(optimizer1, mom)
# zero the parameter gradients
optimizer1.zero_grad()
# forward + backward + optimize
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: ################################################################################check here if broke
# data = tuple(d.cuda() for d in data)
# get the inputs; data is a list of [inputs, labels]
inputs, labels = data
# one cycle policy
if task.upper() == "IMAGENET":
lr, mom = onecycle2.calc()
update_lr(optimizer2, lr)
update_mom(optimizer2, mom)
# 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.eval()
print("Running attack...")
if masked:
adversary = MaskedPGDAttack(
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,
device=device)
else:
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)
shared_x_test = []
shared_y_test = []
adv_untargeted = []
for data in testloader_shared:
x_test = data[0].float().to(device)
y_test = data[1].to(device)
adv_x = adversary.perturb(x_test, y_test)
adv_untargeted.append(adv_x)
shared_x_test.append(x_test)
shared_y_test.append(y_test)
adv_untargeted = torch.cat(adv_untargeted)
shared_x_test = torch.cat(shared_x_test)
shared_y_test = torch.cat(shared_y_test)
# adv_untargeted = adversary.perturb(shared_x_test, shared_y_test)
timestr = time.strftime("%Y%m%d_%H%M%S")
model1 = model1.to("cpu")
model2 = model2.to("cpu")
model1_name = args.save_model_dir + '/{}_{}_{}_model1_{}.pickle'.format(
task, num_shared_classes, percent_shared_data, timestr)
model2_name = args.save_model_dir + '/{}_{}_{}_model2_{}.pickle'.format(
task, num_shared_classes, percent_shared_data, timestr)
adv_name = args.save_model_dir + '/{}_{}_{}_adv_{}.pickle'.format(
task, num_shared_classes, percent_shared_data, timestr)
if savemodel:
print("saving models at", timestr)
torch.save(model1, model1_name)
torch.save(model2, model2_name)
torch.save(adversary, adv_name)
model1_x_test = []
model1_y_test = []
for i in range(len(model1_test_dataset)):
data = model1_test_dataset[i]
model1_x_test.append(data[0])
model1_y_test.append(data[1])
model1_x_test = torch.stack(model1_x_test)
model1_y_test = torch.tensor(model1_y_test)
model2_x_test = []
model2_y_test = []
for i in range(len(model2_test_dataset)):
data = model2_test_dataset[i]
model2_x_test.append(data[0])
model2_y_test.append(data[1])
model1 = model1.to(device)
model2 = model2.to(device)
model2_x_test = torch.stack(model2_x_test)
model2_y_test = torch.tensor(model2_y_test)
model1_x_test = model1_x_test.to(device)
model2_x_test = model2_x_test.to(device)
# 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 run_net(params, transforms):
if params["fixed_seed"] is not None:
torch.backends.cudnn.deterministic = True
random.seed(1)
torch.manual_seed(1)
torch.cuda.manual_seed(1)
np.random.seed(1)
if params["gpu_id"] is not None:
device = torch.device("cuda:{}".format(params["gpu_id"]))
else:
device = torch.device('cpu')
if params["dset"] == "CIFAR10":
concatenated = torch.utils.data.ConcatDataset([
CIFAR10(root='./data', train=True, download=True, transform=None),
CIFAR10(root='./data', train=False, download=True, transform=None)
])
elif params["dset"] == "MNIST":
concatenated = torch.utils.data.ConcatDataset([
MNIST(root='./data', train=True, download=True, transform=None),
MNIST(root='./data', train=False, download=True, transform=None)
])
elif params["dset"] == "FASHIONMNIST":
concatenated = torch.utils.data.ConcatDataset([
FashionMNIST(root='./data',
train=True,
download=True,
transform=None),
FashionMNIST(root='./data',
train=False,
download=True,
transform=None)
])
triplet_test = RandomTripletMiningDataset(concatenated,
train=False,
trans=transforms,
device=None,
params=params)
# Initialize model
if params["dset"] == "CIFAR10":
embedding_net = EmbeddingNet(
in_channels=3,
adjusting_constant=5) # Change this to VGG for CIFAR10
elif params["dset"] == "MNIST" or params["dset"] == "FASHIONMNIST":
embedding_net = EmbeddingNet()
model = TripletNet(embedding_net).to(device)
# Sets the datetime string to use as an identifier for the future
if params["run_id"] is None:
params["run_id"] = '_'.join(
(str(datetime.datetime.now()).split('.')[0].split()))
params["embedding_net_path"] = "arches/embedding_net_{}".format(
params["run_id"])
# Train our model
if params["do_learn"]:
# Initialize loss functions
train_loss = TripletLoss(margin=params["margin"])
test_loader = torch.utils.data.DataLoader(
triplet_test, batch_size=params["batch_size"], shuffle=True
) # Test data is the same as train data but test data is not preshuffled
writer = SummaryWriter(comment="triplet_{0}_{1}_{2}".format(
params["dset"], params["num_pairs"], params["rtm_index"]))
optimizer = optim.Adam(model.parameters(),
lr=params["starting_lr"],
weight_decay=params["weight_decay"])
test(model,
device,
test_loader,
writer,
record_histograms=False,
params=params)
for epoch in range(params["num_epochs"]):
params["curr_epoch"] = epoch
torch.save(model.embedding_net.state_dict(),
params["embedding_net_path"])
triplet_train = RandomTripletMiningDataset(concatenated,
train=True,
trans=transforms,
device=device,
params=params)
train_loader = torch.utils.data.DataLoader(
triplet_train, batch_size=params["batch_size"], shuffle=True)
sample_loader = torch.utils.data.DataLoader(
triplet_train,
batch_size=len(triplet_train) // 100,
shuffle=True
) # Sample our data to get a reference point after every so often
sample_data, sample_targets = next(iter(sample_loader))
if params["show_plots"]:
show_datasets(sample_data)
similar_pair_accuracy = np.round(
len(np.where(sample_targets[0] == sample_targets[1])[0]) /
len(sample_targets[0]), 3)
different_pair_accuracy = np.round(
len(np.where(sample_targets[0] != sample_targets[2])[0]) /
len(sample_targets[0]), 3)
print("Similar pair accuracy:", similar_pair_accuracy)
print("Different pair accuracy:", different_pair_accuracy)
params["different_random_pair_accuracy"] = different_pair_accuracy
params["similar_random_pair_accuracy"] = similar_pair_accuracy
train(model,
device,
train_loader,
train_loss,
epoch,
optimizer,
sample_data,
params=params)
embeddings, targets, indices = test(model,
device,
test_loader,
writer,
epoch=epoch,
params=params)
write_to_tensorboard(
params, writer,
epoch) # Writes to tensorboard at the end of the epoch
writer.export_scalars_to_json(".all_scalars.json")
writer.close()
curr_date = datetime.datetime.now()
date_str = curr_date.strftime("./models/%m_%d_%Y_%H_%M_model_cifar")
print("Saving the full model to: %s" % str(date_str))
torch.save(model, date_str)
print("Model was saved successfully")
loss = F.cross_entropy(out, y_true)
avg_loss.update(loss, img.shape[0])
# Add the labels
y_gt += list(y_true.numpy())
y_pred_label += list(y_pred_label_tmp.numpy())
return avg_loss.avg, y_gt, y_pred_label
if __name__ == "__main__":
trans_img = transforms.Compose([transforms.ToTensor()])
dataset = FashionMNIST("./data/",
train=False,
transform=trans_img,
download=True)
testloader = DataLoader(dataset, batch_size=1024, shuffle=False)
from train_multi_layer import MLP
model_MLP = MLP(10)
model_MLP.load_state_dict(torch.load("./models/MLP.pt"))
from training_conv_net import LeNet
model_conv_net = LeNet(10)
model_conv_net.load_state_dict(torch.load("./models/convNet.pt"))
loss, gt, pred = test(model_MLP, testloader)
with open("multi-layer-net.txt", 'w') as f:
f.write("Loss on Test Data : {}\n".format(loss))
f.write("Accuracy on Test Data : {}\n".format(
def load_dataset(dataset):
train_transform = T.Compose([
T.RandomHorizontalFlip(),
T.RandomCrop(size=32, padding=4),
T.ToTensor(),
T.Normalize(
[0.4914, 0.4822, 0.4465], [0.2023, 0.1994, 0.2010]
) # T.Normalize((0.5071, 0.4867, 0.4408), (0.2675, 0.2565, 0.2761)) # CIFAR-100
])
test_transform = T.Compose([
T.ToTensor(),
T.Normalize(
[0.4914, 0.4822, 0.4465], [0.2023, 0.1994, 0.2010]
) # T.Normalize((0.5071, 0.4867, 0.4408), (0.2675, 0.2565, 0.2761)) # CIFAR-100
])
if dataset == 'cifar10':
data_train = CIFAR10('../cifar10',
train=True,
download=True,
transform=train_transform)
data_unlabeled = MyDataset(dataset, True, test_transform)
data_test = CIFAR10('../cifar10',
train=False,
download=True,
transform=test_transform)
NO_CLASSES = 10
adden = ADDENDUM
no_train = NUM_TRAIN
elif dataset == 'cifar10im':
data_train = CIFAR10('../cifar10',
train=True,
download=True,
transform=train_transform)
#data_unlabeled = CIFAR10('../cifar10', train=True, download=True, transform=test_transform)
targets = np.array(data_train.targets)
#NUM_TRAIN = targets.shape[0]
classes, _ = np.unique(targets, return_counts=True)
nb_classes = len(classes)
imb_class_counts = [500, 5000] * 5
class_idxs = [np.where(targets == i)[0] for i in range(nb_classes)]
imb_class_idx = [
class_id[:class_count]
for class_id, class_count in zip(class_idxs, imb_class_counts)
]
imb_class_idx = np.hstack(imb_class_idx)
no_train = imb_class_idx.shape[0]
# print(NUM_TRAIN)
data_train.targets = targets[imb_class_idx]
data_train.data = data_train.data[imb_class_idx]
data_unlabeled = MyDataset(dataset[:-2], True, test_transform)
data_unlabeled.cifar10.targets = targets[imb_class_idx]
data_unlabeled.cifar10.data = data_unlabeled.cifar10.data[
imb_class_idx]
data_test = CIFAR10('../cifar10',
train=False,
download=True,
transform=test_transform)
NO_CLASSES = 10
adden = ADDENDUM
no_train = NUM_TRAIN
elif dataset == 'cifar100':
data_train = CIFAR100('../cifar100',
train=True,
download=True,
transform=train_transform)
data_unlabeled = MyDataset(dataset, True, test_transform)
data_test = CIFAR100('../cifar100',
train=False,
download=True,
transform=test_transform)
NO_CLASSES = 100
adden = 2000
no_train = NUM_TRAIN
elif dataset == 'fashionmnist':
data_train = FashionMNIST('../fashionMNIST',
train=True,
download=True,
transform=T.Compose([T.ToTensor()]))
data_unlabeled = MyDataset(dataset, True, T.Compose([T.ToTensor()]))
data_test = FashionMNIST('../fashionMNIST',
train=False,
download=True,
transform=T.Compose([T.ToTensor()]))
NO_CLASSES = 10
adden = ADDENDUM
no_train = NUM_TRAIN
elif dataset == 'svhn':
data_train = SVHN('../svhn',
split='train',
download=True,
transform=T.Compose([T.ToTensor()]))
data_unlabeled = MyDataset(dataset, True, T.Compose([T.ToTensor()]))
data_test = SVHN('../svhn',
split='test',
download=True,
transform=T.Compose([T.ToTensor()]))
NO_CLASSES = 10
adden = ADDENDUM
no_train = NUM_TRAIN
return data_train, data_unlabeled, data_test, adden, NO_CLASSES, no_train
from torchvision import transforms
from torchvision.datasets import FashionMNIST
from torch.utils.data import Dataset, DataLoader
DIM_X = 784
### prep dataset and data transforms
# Data augmentation and normalization for training
# Just normalization for validation
data_transforms = {
'train': transforms.Compose([transforms.ToTensor(), lambda x: x.view(-1)]),
'val': transforms.Compose([transforms.ToTensor(), lambda x: x.view(-1)])
}
dataset = FashionMNIST('./FashionMNIST',
download=True,
transform=data_transforms['train'])
data_loader = DataLoader(dataset, batch_size=32)
import seaborn as sns
import copy
import time
import torch
import torch.nn as nn
from torch.optim import Adam
import torch.utils.data as Data
from torchvision import transforms
from torchvision.datasets import FashionMNIST
if __name__ == "__main__":
# 使用fashionMiNist数据，准备训练集
train_data = FashionMNIST(
root="D:/SS/anacondaWork/pytorch/data/FashionMNIST",
train=True,
transform=transforms.ToTensor(),
download=False
)
# 定义一个数据加载器
train_loader = Data.DataLoader(
dataset=train_data,
batch_size=64,
shuffle=False, # 将shuffle设置为false，使得每个批次的训练中样本是固定的，这样有利于在训练中将数据切分为训练集和验证集
num_workers=2
)
print("train_loader的batch数为:", len(train_loader))
# 获取一个批次的数据进行可视化
for step, (b_x, b_y) in enumerate(train_loader):
if step > 0:
ITERS = 30
EVALUATION_CHECKPOINT = 1
AUGMENTATION = False
SESSION_NAME = "sinusoidal_5_100_KP_{}_{}".format(DROPOUT, time.strftime('%Y%m%d-%H%M%S'))
BN_WEIGHT = 0
COV_WEIGHT = 0
CLASSIFIER_TYPE = "dense" # "conv" / "dense"
LOG_DIR = "logs/%s" % SESSION_NAME
EVALUATE_USEFULNESS = True
USEFULNESS_EVAL_SET_SIZE = 1000
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
starttime = time.time()
train_dataset = FashionMNIST('.', train=True, download=True,
transform=transforms.Compose([transforms.ToTensor()]))
test_dataset = FashionMNIST('.', train=False, download=True,
transform=transforms.Compose([transforms.ToTensor()]))
eval_dataset = torch.utils.data.Subset(test_dataset, range(USEFULNESS_EVAL_SET_SIZE))
train_loader = DataLoader(dataset=train_dataset, batch_size=BATCH_SIZE, shuffle=True)
test_loader = DataLoader(dataset=test_dataset, batch_size=BATCH_SIZE, shuffle=False)
eval_loader = DataLoader(dataset=eval_dataset, batch_size=BATCH_SIZE, shuffle=False)
# net = torchnet.FFNet(WIDTH, DEPTH, DROPOUT, OUTPUT_COUNT)
net = torchnet.LeNet()
net.to(device)
print(net)
from torchvision.datasets import FashionMNIST
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
data_train = FashionMNIST("./data/FashionMNIST",
download=True,
train=True,
transform=transforms.Compose([
transforms.Resize((28, 28)),
transforms.ToTensor()
]))
data_val = FashionMNIST("./data/FashionMNIST",
train=False,
download=True,
transform=transforms.Compose([
transforms.Resize((28, 28)),
transforms.ToTensor()
]))
dataloader_train = DataLoader(data_train,
batch_size=1000,
shuffle=True,
num_workers=0)
dataloader_val = DataLoader(data_val, batch_size=1000, num_workers=0)
dataloader_test = DataLoader(data_val,
batch_size=1,
num_workers=0,
shuffle=True)
dataloaders = {
parser = argparse.ArgumentParser(
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('--learning_rate', type=int, default='0.001', help='Learning Rate')
parser.add_argument('--num_epochs', type=int, default=10, help='number of epochs')
parser.add_argument('--batch_size', type=int, default=100, help='Batch Size')
return parser.parse_args()
args = get_args()
normalize = transforms.Normalize(mean=[x/255.0 for x in [125.3, 123.0, 113.9]],
std=[x/255.0 for x in [63.0, 62.1, 66.7]])
transform = transforms.Compose([transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))])
train = FashionMNIST(root='./data', train=True, transform=transform, download=True)
test = FashionMNIST(root='./data', train=False, transform=transform)
train_loader = torch.utils.data.DataLoader(dataset=train,
batch_size=args.batch_size,
shuffle=True)
test_loader = torch.utils.data.DataLoader(dataset=test,
batch_size=args.batch_size,
shuffle=False)
class CNN(nn.Module):
def __init__(self):
super(CNN, self).__init__()
self.layer1 = nn.Sequential(
def test_data():
"""Download and load the test data."""
return FashionMNIST(root=ROOT_DIR_FASHION_MNIST,
download=True,
train=False,
transform=TRANSFORM_FASHION_MNIST)
def main(args):
random = deepcopy(seed)
ts = time.time()
datasets = OrderedDict()
datasets['train'] = FashionMNIST(root='data/fashion',
train=True,
transform=transforms.ToTensor(),
download=True)
tracker_global_train = defaultdict(torch.cuda.FloatTensor)
tracker_global_test = defaultdict(torch.cuda.FloatTensor)
total = np.sum(datasets['train'].targets.data.cpu().numpy() >= 0)
def loss_fn(recon_x, x, q_dist, p_dist):
"""
Variational Auto Encoder loss.
"""
x = x.view(-1)
recon_x = recon_x.view(-1)
rec_loss = torch.nn.functional.binary_cross_entropy(recon_x,
x,
size_average=False)
kl_disc_loss = torch.distributions.kl.kl_divergence(
q_dist, p_dist) # kl_discrete(q_dist, p_dist)
kl_disc_loss = torch.mean(kl_disc_loss, dim=0, keepdim=True)
kl_disc_loss = torch.sum(kl_disc_loss)
return rec_loss + kl_disc_loss, rec_loss, kl_disc_loss
vae = VAE(
encoder_layer_sizes_q=args.encoder_layer_sizes_q,
encoder_layer_sizes_p=args.encoder_layer_sizes_p,
latent_size=args.latent_size,
decoder_layer_sizes=args.decoder_layer_sizes,
conditional=args.conditional,
num_labels=2 if args.conditional else 0 # used to be 10
)
vae = vae.cuda()
# ### # self.latent = DiscreteLatent(self.hyperparams, self.device)
optimizer = torch.optim.Adam(vae.parameters(), lr=args.learning_rate)
for epoch in range(args.epochs):
tracker_epoch = defaultdict(lambda: defaultdict(dict))
for split, dataset in datasets.items():
print("split", split, epoch)
data_loader = DataLoader(dataset=dataset,
batch_size=args.batch_size,
shuffle=split == 'train')
for iteration, (x, y) in enumerate(data_loader):
x = to_var(x)
# convert to tops (0) and bottoms (1)
y[np.logical_or(
y == 0,
np.logical_or(
y == 2,
np.logical_or(y == 3, np.logical_or(y == 4,
y == 6))))] = 0
y[np.logical_or(y == 1, np.logical_or(y == 5, y > 6))] = 1
x = x.view(-1, 784)
y = y.view(-1, 1).cuda()
if args.conditional:
recon_x, alpha_q, alpha_p, alpha_p_sparsemax, alpha_q_lin, alpha_p_lin, z, features = vae(
x, y)
else:
recon_x, alpha_q, alpha_p, alpha_p_sparsemax, alpha_q_lin, alpha_p_lin, z, features = vae(
x)
for i, yi in enumerate(y.data):
id = len(tracker_epoch)
tracker_epoch[id]['x'] = z[i, 0].data
tracker_epoch[id]['y'] = z[i, 1].data
tracker_epoch[id]['label'] = yi[0]
# Form distributions out of alpha_q and alpha_p
q_dist = torch.distributions.one_hot_categorical.OneHotCategorical(
probs=alpha_q)
p_dist = torch.distributions.one_hot_categorical.OneHotCategorical(
probs=alpha_p)
loss, rec, kl = loss_fn(recon_x, x, q_dist, p_dist)
if split == 'train':
optimizer.zero_grad()
loss.backward()
optimizer.step()
tracker_global_train['loss'] = torch.cat(
(tracker_global_train['loss'],
(loss.data / x.size(0)).unsqueeze(-1)))
tracker_global_train['it'] = torch.cat(
(tracker_global_train['it'],
torch.Tensor([epoch * len(data_loader) + iteration
]).cuda()))
# if ((iteration == len(data_loader)-1) and (epoch == args.epochs - 1)):
if iteration % args.print_every == 0 or iteration == len(
data_loader) - 1:
print("Batch %04d/%i, Loss %9.4f" %
(iteration, len(data_loader) - 1, loss.data))
print("recon_x", torch.max(recon_x))
print("recon", rec, "kl", kl)
plt.figure()
plt.figure(figsize=(10, 20))
if args.conditional:
# c= to_var(torch.arange(0,2).repeat(1,10).long().view(-1,1)) # to_var(torch.arange(0,10).long().view(-1,1))
c = to_var(torch.arange(0, 2).long().view(
-1,
1)) # to_var(torch.arange(0,10).long().view(-1,1))
if args.conditional:
# c= to_var(torch.arange(0,2).long().view(-1,1)) # to_var(torch.arange(0,10).long().view(-1,1))
# x, alpha_p, linear_p, features, z = vae.inference(n=c.size(0), c=c)
x, alpha_p, alpha_p_sparsemax, linear_p, features, z = vae.inference(
n=10, c=c)
if 'x' in tracker_global_train.keys():
tracker_global_train['z'] = torch.cat(
(tracker_global_train['z'],
torch.unsqueeze(z, dim=-1)),
dim=-1)
tracker_global_train['x'] = torch.cat(
(tracker_global_train['x'],
torch.unsqueeze(x, dim=-1)),
dim=-1)
tracker_global_train['alpha_p'] = torch.cat(
(tracker_global_train['alpha_p'],
torch.unsqueeze(alpha_p, dim=-1)),
dim=-1)
tracker_global_train[
'alpha_p_sparsemax'] = torch.cat(
(tracker_global_train['alpha_p_sparsemax'],
torch.unsqueeze(alpha_p_sparsemax,
dim=-1)),
dim=-1)
tracker_global_train['weight'] = torch.cat(
(tracker_global_train['weight'],
torch.unsqueeze(linear_p.weight, dim=-1)),
dim=-1)
tracker_global_train['bias'] = torch.cat(
(tracker_global_train['bias'],
torch.unsqueeze(linear_p.bias, dim=-1)),
dim=-1)
tracker_global_train['features'] = torch.cat(
(tracker_global_train['features'],
torch.unsqueeze(features, dim=-1)),
dim=-1)
tracker_global_train['c'] = torch.cat(
(tracker_global_train['c'],
torch.unsqueeze(c, dim=-1)),
dim=-1)
else:
tracker_global_train['z'] = torch.unsqueeze(z,
dim=-1)
tracker_global_train['x'] = torch.unsqueeze(x,
dim=-1)
tracker_global_train['alpha_p'] = torch.unsqueeze(
alpha_p, dim=-1)
tracker_global_train[
'alpha_p_sparsemax'] = torch.unsqueeze(
alpha_p_sparsemax, dim=-1)
tracker_global_train['weight'] = torch.unsqueeze(
linear_p.weight, dim=-1)
tracker_global_train['bias'] = torch.unsqueeze(
linear_p.bias, dim=-1)
tracker_global_train['features'] = torch.unsqueeze(
features, dim=-1)
tracker_global_train['c'] = torch.unsqueeze(c,
dim=-1)
else:
x, alpha_p, alpha_p_sparsemax, linear_p, features = vae.inference(
n=c.size(0))
folder_name = str(ts) + "_fashion_Adam_random_" + str(
random)
if not os.path.exists(
os.path.join(args.fig_root, folder_name)):
if not (os.path.exists(os.path.join(args.fig_root))):
os.mkdir(os.path.join(args.fig_root))
os.mkdir(os.path.join(args.fig_root, folder_name))
z_folder = os.path.join(
folder_name, "epoch_%i_iter_%i/" % (epoch, iteration))
if not os.path.exists(os.path.join(args.fig_root,
z_folder)):
if not (os.path.exists(os.path.join(args.fig_root))):
os.mkdir(os.path.join(args.fig_root))
os.mkdir(os.path.join(args.fig_root, z_folder))
#for p in range(c.size(0)):
for p in range(10):
plt.clf()
plt.close()
plt.imshow(x[p].view(28, 28).data.cpu().numpy())
plt.axis('off')
plt.savefig(os.path.join(args.fig_root, z_folder,
"%i.png" % (p)),
dpi=300)
plt.clf()
plt.close()
# Plot losses
plt.plot(tracker_global_train['it'].data.cpu().numpy(),
tracker_global_train['loss'].data.cpu().numpy())
plt.xlabel('Iteration')
plt.ylabel('Loss')
plt.ylim(0, 550)
plt.savefig(os.path.join(args.fig_root, folder_name, "loss.png"))
plt.clf()
plt.close()
# Save data
torch.save(tracker_global_train,
'tracker_fashion_Adam_random_' + str(random) + '.pt')
train_losses.append(running_loss / i) # 1 epoch loss
train_acc.append(n_acc / n) # 모든 batch에서 계산한 정확도를 1 epoch때 나누어서 계산
val_acc.append(eval_net(net, test_loader, device))
print(epoch, train_losses[-1], train_acc[-1], val_acc[-1], flush=True)
# for epoch 끝
if __name__ == '__main__':
## 학습용 데이터 받기
myPath = "/media/jsh/CA02176502175633/Users/Desktop/Documents/data"
fashion_mnist_train = FashionMNIST(myPath + "/FashionMNIST",
train=True,
download=True,
transform=transforms.ToTensor())
# 테스트용 데이터 받기
# Tensor형태로 받는다. ==> dataset로 만들어 버리기 위함.
fashion_mnist_test = FashionMNIST(myPath + "/FashionMNIST",
train=False,
download=True,
transform=transforms.ToTensor())
batch_size = 128 # batch size 설정
## 데이터를 batch 단위로 묶어놓는다.
train_loader = DataLoader(fashion_mnist_train,
batch_size=batch_size,
shuffle=True)
net.features[0] = nn.Conv2d(1,
64,
kernel_size=7,
stride=2,
padding=3,
bias=False)
batch_size = 32
transform = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize((0.5), (0.5))
])
fashion_mnist_trainval = FashionMNIST("FashionMNIST",
train=True,
download=True,
transform=transform)
fashion_mnist_test = FashionMNIST("FashionMNIST",
train=False,
download=True,
transform=transform)
n_samples = len(fashion_mnist_trainval)
train_size = int(len(fashion_mnist_trainval) * 0.8)
val_size = n_samples - train_size
train_dataset, val_dataset = torch.utils.data.random_split(
fashion_mnist_trainval, [train_size, val_size])
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=True)
images = Variable(images, volatile=True)
if cuda:
images = images.cuda()
embeddings[k:k + len(images)] = model.get_embedding(
images).data.cpu().numpy()
labels[k:k + len(images)] = target.numpy()
k += len(images)
return embeddings, labels
mean, std = 0.28604059698879553, 0.35302424451492237
train_dataset = FashionMNIST('../data/FashionMNIST',
train=True,
download=True,
transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((mean, ), (std, ))
]))
test_dataset = FashionMNIST('../data/FashionMNIST',
train=False,
download=True,
transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((mean, ), (std, ))
]))
n_classes = 10
batch_size = 256
kwargs = {'num_workers': 1, 'pin_memory': True} if cuda else {}
train_loader = torch.utils.data.DataLoader(train_dataset,
new_classifier = nn.Sequential(*list(model.classifier.children())[:-1])
model.classifier = new_classifier
# In[17]:
train_transforms = transforms.Compose([
# transforms.RandomHorizontalFlip(),
# utils.RandomRotation(),
# utils.RandomTranslation(),
# utils.RandomVerticalFlip(),
transforms.ToTensor()
# transforms.Normalize((0.1307,), (0.3081,))
]
)
trainset = FashionMNIST('fashion-mnist/data/', train=True, download=True, transform=train_transforms)
train_loader = DataLoader(trainset, batch_size=64)
# In[18]:
def extractFeatures(net, loader):
net.eval()
size = loader.dataset.train_data.shape[0]
batch_size = loader.batch_size
num_batches = size // batch_size
num_batches = num_batches if loader.drop_last else num_batches + 1
features = []
from torchvision.models import vgg16
from torchvision import transforms
# ## Formla to claclulate output of conv layer
#
# input is : n,w,h
# output :n_new , w_new , h_new
# n_new = out_channels
# w_new = (w - kernal + 2 padding)/stride + 1
# h_new = (h - kernal + 2 padding)/stride + 1
# In[5]:
batch = 20
train_data = FashionMNIST(r'C:\Users\elfakhrany\Documents\session2\data',
train=True,
transform=transforms.ToTensor(),
download=False)
trainloader = DataLoader(train_data, batch_size=batch, shuffle=True)
# In[8]:
# let's get some Info About data
iterator = iter(trainloader)
images, labels = iterator.next()
print("Num of images in training data is {}".format(len(train_data)))
print('Images shape is : {}'.format(images.shape))
print('Lables shape is : {}'.format(labels.shape))
# In[9]:
#check labels
def get_loader(args, is_train):
dataset = args.dataset
if dataset.lower() == 'cifar10':
dset = get_cifar(is_train)
elif dataset == 'GaussianLine':
# dset = GaussianLine(args.fdata)
xdir = np.load(args.fxdir) if args.fxdir else None
dset = GaussianLine(args.d, bt=args.bt, dlen=args.dlen, xdir=xdir)
elif dataset == 'GaussianMixture':
dset = GaussianMixture(dlen=args.dlen)
elif dataset == 'uniform':
dset = Uniform(dlen=args.dlen)
elif dataset == 'GAS8':
dset = GAS8(args.norm_by_col)
elif dataset == 'GAS16':
dset = GAS16(args.norm_by_col)
elif dataset == 'GAS128':
dset = GAS128(args.norm_by_col)
elif dataset == 'gas':
dset = gas(args.n_pts)
elif dataset == 'hepmass':
dset = HEPMASS()
elif dataset == 'miniboone':
dset = MINIBooNE()
elif dataset == 'MNISTtab':
# treat mnist as tabular data
dset = MNISTtab(args.pca_dim)
elif dataset == 'MNIST':
channel = 1
image_size = 28
lambd = 1e-5
transform = transforms.Compose(
[transforms.Resize(image_size),
transforms.ToTensor()])
dset = MNIST(os.path.join('datasets', 'mnist'),
train=is_train,
download=True,
transform=transform)
elif dataset == 'FMNIST':
channel = 1
image_size = 28
lambd = 1e-6
transform = transforms.Compose(
[transforms.Resize(image_size),
transforms.ToTensor()])
dset = FashionMNIST(os.path.join('datasets', 'fmnist'),
train=is_train,
download=True,
transform=transform)
else:
print('args.dataset:', dataset)
raise NotImplementedError(
'Sorry but we are not supporting dataset {}. \n Ciao~! :)'.format(
dataset))
if is_train and args.use_val:
n_train = int(0.8 * len(dset))
n_val = len(dset) - n_train
dset_train, dset_val = torch.utils.data.random_split(
dset, [n_train, n_val])
train_loader = data.DataLoader(dset_train,
args.bt,
shuffle=True,
num_workers=args.num_workers)
val_loader = data.DataLoader(dset_val,
args.bt_test,
shuffle=False,
num_workers=args.num_workers)
return train_loader, val_loader
return data.DataLoader(dset, args.bt, shuffle=is_train)
# data loading and transforming
from torchvision.datasets import FashionMNIST
from torch.utils.data import DataLoader
from torchvision import transforms
import numpy as np
# The output of torchvision datasets are PILImage images of range [0, 1].
# We transform them to Tensors for input into a CNN
## Define a transform to read the data in as a tensor
data_transform = transforms.ToTensor()
# choose the training and test datasets
train_data = FashionMNIST(root='data',
train=True,
download=True,
transform=data_transform)
test_data = FashionMNIST(root='data',
train=False,
download=True,
transform=data_transform)
# Print out some stats about the training and test data
print('Train data, number of images: ', len(train_data))
print('Test data, number of images: ', len(test_data))
# prepare data loaders, set the batch_size
batch_size = 20
train_loader = DataLoader(train_data, batch_size=batch_size, shuffle=True)
def load_data(config):
normal_class = config['normal_class']
batch_size = config['batch_size']
img_size = config['image_size']
if config['dataset_name'] in ['cifar10']:
img_transform = transforms.Compose([
transforms.ToTensor(),
])
os.makedirs("./Dataset/CIFAR10/train", exist_ok=True)
dataset = CIFAR10('./Dataset/CIFAR10/train', train=True, download=True, transform=img_transform)
dataset.data = dataset.data[np.array(dataset.targets) == normal_class]
dataset.targets = [normal_class] * dataset.data.shape[0]
train_set, val_set = torch.utils.data.random_split(dataset, [dataset.data.shape[0] - 851, 851])
os.makedirs("./Dataset/CIFAR10/test", exist_ok=True)
test_set = CIFAR10("./Dataset/CIFAR10/test", train=False, download=True, transform=img_transform)
elif config['dataset_name'] in ['mnist']:
img_transform = transforms.Compose([
transforms.Resize(img_size),
transforms.ToTensor(),
])
os.makedirs("./Dataset/MNIST/train", exist_ok=True)
dataset = MNIST('./Dataset/MNIST/train', train=True, download=True, transform=img_transform)
dataset.data = dataset.data[np.array(dataset.targets) == normal_class]
dataset.targets = [normal_class] * dataset.data.shape[0]
train_set, val_set = torch.utils.data.random_split(dataset, [dataset.data.shape[0] - 851, 851])
os.makedirs("./Dataset/MNIST/test", exist_ok=True)
test_set = MNIST("./Dataset/MNIST/test", train=False, download=True, transform=img_transform)
elif config['dataset_name'] in ['fashionmnist']:
img_transform = transforms.Compose([
transforms.Resize(img_size),
transforms.ToTensor(),
])
os.makedirs("./Dataset/FashionMNIST/train", exist_ok=True)
dataset = FashionMNIST('./Dataset/FashionMNIST/train', train=True, download=True, transform=img_transform)
dataset.data = dataset.data[np.array(dataset.targets) == normal_class]
dataset.targets = [normal_class] * dataset.data.shape[0]
train_set, val_set = torch.utils.data.random_split(dataset, [dataset.data.shape[0] - 851, 851])
os.makedirs("./Dataset/FashionMNIST/test", exist_ok=True)
test_set = FashionMNIST("./Dataset/FashionMNIST/test", train=False, download=True, transform=img_transform)
elif config['dataset_name'] in ['brain_tumor', 'head_ct']:
img_transform = transforms.Compose([
transforms.Resize([img_size, img_size]),
transforms.Grayscale(num_output_channels=1),
transforms.ToTensor()
])
root_path = 'Dataset/medical/' + config['dataset_name']
train_data_path = root_path + '/train'
test_data_path = root_path + '/test'
dataset = ImageFolder(root=train_data_path, transform=img_transform)
load_dataset = DataLoader(dataset, batch_size=batch_size, shuffle=True)
train_dataset_array = next(iter(load_dataset))[0]
my_dataset = TensorDataset(train_dataset_array)
train_set, val_set = torch.utils.data.random_split(my_dataset, [train_dataset_array.shape[0] - 5, 5])
test_set = ImageFolder(root=test_data_path, transform=img_transform)
elif config['dataset_name'] in ['coil100']:
img_transform = transforms.Compose([
transforms.ToTensor()
])
root_path = 'Dataset/coil100/' + config['dataset_name']
train_data_path = root_path + '/train'
test_data_path = root_path + '/test'
dataset = ImageFolder(root=train_data_path, transform=img_transform)
load_dataset = DataLoader(dataset, batch_size=batch_size, shuffle=True)
train_dataset_array = next(iter(load_dataset))[0]
my_dataset = TensorDataset(train_dataset_array)
train_set, val_set = torch.utils.data.random_split(my_dataset, [train_dataset_array.shape[0] - 5, 5])
test_set = ImageFolder(root=test_data_path, transform=img_transform)
elif config['dataset_name'] in ['MVTec']:
data_path = 'Dataset/MVTec/' + normal_class + '/train'
data_list = []
orig_transform = transforms.Compose([
transforms.Resize(img_size),
transforms.ToTensor()
])
orig_dataset = ImageFolder(root=data_path, transform=orig_transform)
train_orig, val_set = torch.utils.data.random_split(orig_dataset, [len(orig_dataset) - 25, 25])
data_list.append(train_orig)
for i in range(3):
img_transform = transforms.Compose([
transforms.Resize(img_size),
transforms.RandomAffine(0, scale=(1.05, 1.2)),
transforms.ToTensor()])
dataset = ImageFolder(root=data_path, transform=img_transform)
data_list.append(dataset)
dataset = ConcatDataset(data_list)
train_loader = torch.utils.data.DataLoader(dataset, batch_size=800, shuffle=True)
train_dataset_array = next(iter(train_loader))[0]
train_set = TensorDataset(train_dataset_array)
test_data_path = 'Dataset/MVTec/' + normal_class + '/test'
test_set = ImageFolder(root=test_data_path, transform=orig_transform)
train_dataloader = torch.utils.data.DataLoader(
train_set,
batch_size=batch_size,
shuffle=True,
)
val_dataloader = torch.utils.data.DataLoader(
val_set,
batch_size=batch_size,
shuffle=True,
)
test_dataloader = torch.utils.data.DataLoader(
test_set,
batch_size=batch_size,
shuffle=True,
)
return train_dataloader, val_dataloader, test_dataloader
#########################################################################
# Name of model file
#########################################################################
model_name = args.model_name + '.pt'
#########################################################################
# Training
#########################################################################
if args.mode == 'train':
#load dataset to loader
dataset = FashionMNIST("./data",
train=True,
download=True,
transform=transform)
data_loader = torch.utils.data.DataLoader(dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.number_worker)
print(len(dataset))
net = SimpleCNN().cuda()
# number of epochs
def weights_init(m):
classname = m.__class__.__name__
if classname.find('Conv') != -1:
torch.nn.init.xavier_uniform(m.weight)
from torchvision.datasets import FashionMNIST
from torchvision import transforms
# get the training data
train_data = FashionMNIST(
"data/FashionMNIST",
train=True,
download=True,
transform=transforms.ToTensor() # このままだとPIL形式なので、これをRGBの順の形式にかえる
)
# make the test data
test_data = FashionMNIST("data/FashionMNIST",
train=False,
download=True,
transform=transforms.ToTensor())
from torch.utils.data import TensorDataset, DataLoader
# make the data loader
batch_size = 128
train_loader = DataLoader(train_data, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(test_data, batch_size=batch_size, shuffle=False)
# CNN の構築
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import optim
model.load_state_dict(checkpoint)
# Test 1
data = next(iter(dataloader))
image = data[0][0][0].numpy()
label = data[1][0].numpy()
plt.imshow(image)
plt.show()
pred = model(Variable(torch.from_numpy(image.reshape(1, -1))))
pred = pred.view(28, 28).detach().numpy()
plt.imshow(pred)
plt.show()
# Test 2
from torchvision.datasets import FashionMNIST
dataset = FashionMNIST('../data', transform=img_transform, download=True)
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
data = next(iter(dataloader))
image = data[0][0][0].numpy()
label = data[1][0].numpy()
plt.imshow(image)
plt.show()
pred = model(Variable(torch.from_numpy(image.reshape(1, -1))))
pred = pred.view(28, 28).detach().numpy()
plt.imshow(pred)
plt.show()
# Conclusion :
# model performs good on digital dataset - MNIST
# model performs bad on FashionMNIST
# So this simple autoencoder seems not good for generalization
# Update the weights
model.optimizer.step()
return avg_loss.avg, avg_accuracy
if __name__ == "__main__":
device = torch.device('cpu')
model = MLNN(10).to(device)
model.optimizer = model.get_optimizer(opt_name, lr, l2_penalty, momentum)
trans_img = transforms.Compose([transforms.ToTensor()])
traindatasetinital = FashionMNIST("./data/", train=True, transform=trans_img, download=True)
tracktrainloss = []
tracktrainaccuracy = []
trackvalidationloss = []
trackvalidationaccuracy = []
for i in range(epochs):
traindataset, validationdataset = random_split(traindatasetinital, [50000,10000])
trainloader = DataLoader(traindataset, batch_size=1024, shuffle=True)
validationloader = DataLoader(validationdataset, batch_size=1024, shuffle=True)
trainloss, trainaccuracy = train1epoch(model, trainloader, device)
tracktrainloss.append(trainloss)
tracktrainaccuracy.append(trainaccuracy)
validationloss, validationaccuracy = train1epoch(model, validationloader, device)
trackvalidationloss.append(validationloss)
trackvalidationaccuracy.append(validationaccuracy)
def prepare_data(self):
"""
Saves FashionMNIST files to data_dir
"""
FashionMNIST(self.data_dir, train=True, download=True, transform=transform_lib.ToTensor())
FashionMNIST(self.data_dir, train=False, download=True, transform=transform_lib.ToTensor())
def init_dataloaders(args):
if args.dataset == 'omniglot':
from Data.omniglot import Omniglot
from torchvision.datasets import MNIST, FashionMNIST
args.is_classification_task = True
args.prob_train, args.prob_test, args.prob_ood = 0.5, 0.25, 0.25
args.n_train_cls = 900
args.n_val_cls = 100
args.n_train_samples = 10
args.input_size = [1, 28, 28]
Omniglot_dataset = Omniglot(args.folder).data
Omniglot_dataset = torch.from_numpy(Omniglot_dataset).type(
torch.float).to(args.device)
meta_train_dataset = Omniglot_dataset[:args.n_train_cls]
meta_train_train = meta_train_dataset[:, :args.n_train_samples, :, :]
meta_train_test = meta_train_dataset[:, args.n_train_samples:, :, :]
meta_val_dataset = Omniglot_dataset[args.n_train_cls:(
args.n_train_cls + args.n_val_cls)]
#TODO(figure out the bug when there is only a single class fed to the valid dataloader)
meta_val_train = meta_val_dataset[:, :args.n_train_samples, :, :]
meta_val_test = meta_val_dataset[:, args.n_train_samples:, :, :]
cl_dataset = Omniglot_dataset
cl_ood_dataset1 = MNIST(args.folder, train=True, download=True)
cl_ood_dataset2 = FashionMNIST(args.folder, train=True, download=True)
cl_ood_dataset1, _ = order_and_split(cl_ood_dataset1.data,
cl_ood_dataset1.targets)
cl_ood_dataset2, _ = order_and_split(cl_ood_dataset2.data,
cl_ood_dataset2.targets)
cl_ood_dataset1 = cl_ood_dataset1[:, :, None, :, :]
cl_ood_dataset2 = cl_ood_dataset2[:, :, None, :, :]
cl_ood_dataset1 = cl_ood_dataset1.type(torch.float).to(args.device)
cl_ood_dataset2 = cl_ood_dataset2.type(torch.float).to(args.device)
elif args.dataset == "tiered-imagenet":
from Data.tiered_imagenet import NonEpisodicTieredImagenet
args.prob_train, args.prob_test, args.prob_ood = 0.3, 0.3, 0.4
args.is_classification_task = True
args.n_train_cls = 100
args.n_val_cls = 100
args.n_train_samples = 500
args.input_size = [3, 64, 64]
tiered_dataset = NonEpisodicTieredImagenet(args.folder, split="train")
meta_train_dataset = tiered_dataset.data[:args.n_train_cls]
meta_train_train = meta_train_dataset[:, :args.n_train_samples, ...]
meta_train_test = meta_train_dataset[:, args.n_train_samples:, ...]
meta_val_dataset = tiered_dataset.data[args.n_train_cls:(
args.n_train_cls + args.n_val_cls)]
meta_val_train = meta_val_dataset[:, :args.n_train_samples, :, :]
meta_val_test = meta_val_dataset[:, args.n_train_samples:, :, :]
cl_dataset = tiered_dataset.data
set_trace()
cl_ood_dataset1 = tiered_dataset.data[(args.n_train_cls +
args.n_val_cls):]
## last results computed with this split
#cl_ood_dataset1 = tiered_dataset.data[200:300]
cl_ood_dataset2 = NonEpisodicTieredImagenet(args.folder,
split="val").data
#cl_dataset = cl_dataset.type(torch.float)#.to(args.device)
cl_ood_dataset1 = cl_ood_dataset1.type(torch.float) #.to(args.device)
cl_ood_dataset2 = cl_ood_dataset2.type(torch.float) #.to(args.device)
elif args.dataset == "harmonics":
'''under construction'''
from data.harmonics import Harmonics
args.is_classification_task = False
args.input_size = [1]
def make_dataset(train: bool = True) -> torch.Tensor:
return torch.from_numpy(Harmonics(train=train).data).float()
dataset = make_dataset()
meta_train_dataset = dataset[:500]
meta_train_train = meta_train_dataset[:, :40]
meta_train_test = meta_train_dataset[:, 40:]
meta_val_dataset = dataset[500:]
meta_val_train = meta_val_dataset[:, :40]
meta_val_test = meta_val_dataset[:, 40:]
if args.mode == 'train':
cl_dataset = dataset
cl_ood_dataset1 = make_dataset(train=False)
cl_ood_dataset2 = make_dataset(train=False)
cl_ood_dataset3 = make_dataset(train=False)
args.prob_train, args.prob_test, args.prob_ood = 0.6, 0., 0.4
else:
raise NotImplementedError('Unknown dataset `{0}`.'.format(
args.dataset))
meta_train_dataloader = MetaDataset(meta_train_train,
meta_train_test,
args=args,
n_shots_tr=args.num_shots,
n_shots_te=args.num_shots_test,
n_way=args.num_ways)
meta_val_dataloader = MetaDataset(meta_val_train,
meta_val_test,
args=args,
n_shots_tr=args.num_shots,
n_shots_te=args.num_shots_test,
n_way=args.num_ways)
meta_train_dataloader = torch.utils.data.DataLoader(
meta_train_dataloader, batch_size=args.batch_size)
meta_val_dataloader = torch.utils.data.DataLoader(
meta_val_dataloader, batch_size=args.batch_size)
cl_dataloader = StreamDataset(cl_dataset,
cl_ood_dataset1,
cl_ood_dataset2,
n_shots=args.num_shots,
n_way=args.num_ways,
prob_statio=args.prob_statio,
prob_train=args.prob_train,
prob_test=args.prob_test,
prob_ood=args.prob_ood,
args=args)
cl_dataloader = torch.utils.data.DataLoader(cl_dataloader, batch_size=1)
del meta_train_dataset, meta_train_train, meta_train_test, meta_val_dataset,\
meta_val_train, meta_val_test, cl_dataset, cl_ood_dataset1, cl_ood_dataset2
return meta_train_dataloader, meta_val_dataloader, cl_dataloader
def test(self):
states = torch.load(os.path.join(self.args.log, 'checkpoint.pth'),
map_location=self.config.device)
score = RefineNetDilated(self.config).to(self.config.device)
score = torch.nn.DataParallel(score)
score.load_state_dict(states[0])
if not os.path.exists(self.args.image_folder):
os.makedirs(self.args.image_folder)
score.eval()
if self.config.data.dataset == 'MNIST' or self.config.data.dataset == 'FashionMNIST':
transform = transforms.Compose([
transforms.Resize(self.config.data.image_size),
transforms.ToTensor()
])
if self.config.data.dataset == 'MNIST':
dataset = MNIST(os.path.join(self.args.run, 'datasets',
'mnist'),
train=True,
download=True,
transform=transform)
else:
dataset = FashionMNIST(os.path.join(self.args.run, 'datasets',
'fmnist'),
train=True,
download=True,
transform=transform)
dataloader = DataLoader(dataset,
batch_size=100,
shuffle=True,
num_workers=4)
data_iter = iter(dataloader)
samples, _ = next(data_iter)
samples = samples.cuda()
samples = torch.rand_like(samples)
all_samples = self.Langevin_dynamics(samples, score, 1000, 0.00002)
for i, sample in enumerate(tqdm.tqdm(all_samples)):
sample = sample.view(100, self.config.data.channels,
self.config.data.image_size,
self.config.data.image_size)
if self.config.data.logit_transform:
sample = torch.sigmoid(sample)
torch.save(
sample,
os.path.join(self.args.image_folder,
'samples_{}.pth'.format(i)))
elif self.config.data.dataset == 'CELEBA':
dataset = CelebA(
root=os.path.join(self.args.run, 'datasets', 'celeba'),
split='test',
transform=transforms.Compose([
transforms.CenterCrop(140),
transforms.Resize(self.config.data.image_size),
transforms.ToTensor(),
]),
download=True)
dataloader = DataLoader(dataset,
batch_size=64,
shuffle=True,
num_workers=4)
samples, _ = next(iter(dataloader))
samples = torch.rand(100,
3,
self.config.data.image_size,
self.config.data.image_size,
device=self.config.device)
all_samples = self.Langevin_dynamics(samples, score, 1000, 0.00002)
for i, sample in enumerate(tqdm.tqdm(all_samples)):
sample = sample.view(100, self.config.data.channels,
self.config.data.image_size,
self.config.data.image_size)
if self.config.data.logit_transform:
sample = torch.sigmoid(sample)
torch.save(
sample,
os.path.join(self.args.image_folder,
'samples_{}.pth'.format(i)))
else:
transform = transforms.Compose([
transforms.Resize(self.config.data.image_size),
transforms.ToTensor()
])
if self.config.data.dataset == 'CIFAR10':
dataset = CIFAR10(os.path.join(self.args.run, 'datasets',
'cifar10'),
train=True,
download=True,
transform=transform)
dataloader = DataLoader(dataset,
batch_size=100,
shuffle=True,
num_workers=4)
data_iter = iter(dataloader)
samples, _ = next(data_iter)
samples = samples.cuda()
samples = torch.rand_like(samples)
all_samples = self.Langevin_dynamics(samples, score, 1000, 0.00002)
for i, sample in enumerate(tqdm.tqdm(all_samples)):
sample = sample.view(100, self.config.data.channels,
self.config.data.image_size,
self.config.data.image_size)
if self.config.data.logit_transform:
sample = torch.sigmoid(sample)
torch.save(
sample,
os.path.join(self.args.image_folder,
'samples_{}.pth'.format(i)))
def get_data(dataset_used, batch_size, get_mean_std=False):
if dataset_used == 'CIFAR10':
data = CIFAR10('datasets/',
train=True,
transform=transform,
download=True)
elif dataset_used == 'MNIST':
data = MNIST('datasets/',
train=True,
transform=original_transform,
download=True)
elif dataset_used == 'FashionMNIST':
data = FashionMNIST('datasets/',
train=True,
transform=gray_scale_transform,
download=True)
elif dataset_used == 'MPII':
data = MPIIDataSet('datasets/mpii_human_pose_v1')
elif dataset_used == 'UTD':
data = ImageFolder('datasets/UTD-MHAD/Image/', transform=transform)
elif dataset_used == 'UTDVideo':
data = UTDVideo('datasets/UTD-MHAD/Depth/')
if get_mean_std:
mean = torch.zeros(3)
std = torch.zeros(3)
print('Computing mean and std...')
full_data_loader = DataLoader(data,
batch_size=batch_size,
shuffle=False,
num_workers=os.cpu_count())
for idx, batch in enumerate(full_data_loader):
if dataset_used == 'MNIST':
img, mask = batch[0], None
elif dataset_used == 'MPII':
img, mask = batch['image'], batch['mask']
for i in range(3):
mean[i] += img[:, i, :, :].mean()
std[i] += img[:, i, :, :].std()
mean.div_(idx)
std.div_(idx)
print(mean, std)
data_size = len(data)
train_size = data_size * 9 // 10
print("Train size: ", train_size)
val_size = data_size - train_size
train_set, val_set = random_split(data, [train_size, val_size])
train_collate_func = None if session_num == 0 else train_collate
val_collate_func = None if session_num == 0 else val_collate
train_data_loader = DataLoader(train_set,
batch_size=batch_size,
shuffle=True,
num_workers=os.cpu_count(),
collate_fn=train_collate_func)
val_data_loader = DataLoader(val_set,
batch_size=1,
shuffle=True,
num_workers=0,
collate_fn=val_collate_func)
return train_data_loader, val_data_loader, train_size
def test_fashion_mnist_doesnt_load_mnist(self):
MNIST(root=self.test_dir, download=True)
FashionMNIST(root=self.test_dir, download=True)
