def prepare_data(args, dir_path, shuffle=True):
"""prepare data for training and testing"""
kwargs = {'num_workers': 1, 'pin_memory': True} if args.cuda else {}
image_size = global_conf['image_size']
resize = None
if len(image_size) == 2:
resize = image_size
elif len(image_size) > 2:
resize = (image_size[-2], image_size[-1])
else:
raise RuntimeError("config image size is illegal")
train_loader = prepare_data_mnist(args.batch_size,
dir_path,
train=True,
shuffle=shuffle,
resize=resize,
**kwargs)
test_loader = prepare_data_mnist(args.batch_size,
dir_path,
train=False,
shuffle=shuffle,
resize=resize,
**kwargs)
return train_loader, test_loader
def prepare_data(args, dir_path, shuffle=True):
"""prepare data for training and testing"""
kwargs = {'num_workers': 1, 'pin_memory': True} if args.cuda else {}
train_loader = prepare_data_mnist(
args.batch_size, dir_path, train=True, shuffle=shuffle, **kwargs)
test_loader = prepare_data_mnist(
args.batch_size, dir_path, train=False, shuffle=shuffle, **kwargs)
return train_loader, test_loader
def prepare_data(args, dir_path, shuffle=True):
"""prepare data for training/testing"""
kwargs = {'num_workers': 1, 'pin_memory': True} if args.cuda else {}
train_loader1 = prepare_data_mnist(
args.batch_size, dir_path, train=True, shuffle=shuffle, **kwargs)
train_loader2 = DataLoader(
train_loader1.dataset, batch_size=args.batch_size, shuffle=True, **kwargs)
test_loader1 = prepare_data_mnist(
args.batch_size, dir_path, train=False, shuffle=shuffle, **kwargs)
test_loader2 = DataLoader(
test_loader1.dataset, batch_size=args.batch_size, shuffle=True, **kwargs)
return [train_loader1, train_loader2], [test_loader1, test_loader2]
if __name__ == '__main__':
test_conf = set_test_config()
checkpoint_path = test_conf['checkpoint_path']
save_dir = test_conf['save_dir']
traversal_vals = torch.tensor(test_conf['traversal_vals'])
model_checkpoint = torch.load(checkpoint_path)
train_args = model_checkpoint['train_args']
train_conf = model_checkpoint['train_config']
model_state_dict = model_checkpoint['state_dict']
img_size = train_conf['image_size']
model = SparseVAE(img_size[0]*img_size[1], train_args.hidden_sizes,
train_args.dim_z, train_args.alpha)
model.c = model_checkpoint['extra']['model_c'] # update c
model.load_state_dict(model_state_dict)
# mnist dataset
test_loader = prepare_data_mnist(train_args.batch_size, train_conf['data_dir'],
train=False, shuffle=False)
images, _ = iter(test_loader).next()
# traversal
print(traversal_vals)
img_size = (1, *img_size)
for i in range(5):
interpretability_traversal(model, traversal_vals, img_size,
os.path.join(save_dir, "test{}.gif").format(i), images[i])