def run(count, memory_size, iteration, device='cuda'): traintransform = transforms.Compose( [transforms.RandomRotation(20), transforms.ToTensor()]) trainset = torchvision.datasets.MNIST(root='./data/mnist', train=True, download=True, transform=traintransform) trainloader = torch.utils.data.DataLoader(trainset, pin_memory=True, batch_size=128, shuffle=True, num_workers=10) testtransform = transforms.Compose([transforms.ToTensor()]) testset = torchvision.datasets.MNIST(root='./data/mnist', train=False, download=True, transform=testtransform) testloader = torch.utils.data.DataLoader(testset, pin_memory=True, batch_size=128, shuffle=True, num_workers=10) base_dir = os.path.join('mnist_' + str(memory_size), "4") model = MnistDraw(count, memory_size) optimizer = optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=1e-3) from datetime import datetime current_time = datetime.now().strftime('%b%d_%H-%M-%S') trial = Trial( model, optimizer, nn.MSELoss(reduction='sum'), ['loss'], pass_state=True, callbacks=[ tm.kl_divergence(MU, LOGVAR), callbacks.MostRecent( os.path.join(base_dir, 'iter_' + str(iteration) + '.{epoch:02d}.pt')), callbacks.GradientClipping(5), callbacks.ExponentialLR(0.99), callbacks.TensorBoardImages(comment=current_time, name='Prediction', write_each_epoch=True, key=torchbearer.Y_PRED), callbacks.TensorBoardImages(comment=current_time + '_mnist', name='Target', write_each_epoch=True, key=torchbearer.Y_TRUE) ]).with_generators(train_generator=trainloader, val_generator=testloader).to(device) trial.run(100)
def draw(count, memory_size, file, device='cuda'): testtransform = transforms.Compose([transforms.ToTensor()]) testset = torchvision.datasets.MNIST(root='./data/mnist', train=False, download=True, transform=testtransform) testloader = torch.utils.data.DataLoader(testset, pin_memory=True, batch_size=128, shuffle=True, num_workers=10) base_dir = os.path.join('mnist_' + str(memory_size), "6") model = MnistDraw(count, memory_size, output_stages=True) optimizer = optim.SGD(filter(lambda p: p.requires_grad, model.parameters()), lr=0) from datetime import datetime current_time = datetime.now().strftime('%b%d_%H-%M-%S') from visualise import StagesGrid trial = Trial(model, optimizer, nn.MSELoss(reduction='sum'), ['loss'], pass_state=True, callbacks=[ callbacks.TensorBoardImages(comment=current_time, nrow=10, num_images=20, name='Prediction', write_each_epoch=True, key=torchbearer.Y_PRED, pad_value=1), callbacks.TensorBoardImages(comment=current_time + '_mnist', nrow=10, num_images=20, name='Target', write_each_epoch=False, key=torchbearer.Y_TRUE, pad_value=1), StagesGrid('mnist_stages.png', STAGES, 20) ]).load_state_dict(torch.load(os.path.join(base_dir, file)), resume=False).with_generators(train_generator=testloader, val_generator=testloader).for_train_steps(1).for_val_steps(1).to(device) trial.run() # Evaluate doesn't work with tensorboard in torchbearer, seems to have been fixed in most recent version
def draw(file, device='cuda'): transform_test = transforms.Compose([transforms.ToTensor()]) testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform_test) testloader = torch.utils.data.DataLoader(testset, batch_size=100, shuffle=False, num_workers=10) base_dir = 'cifar_vae' model = CifarVAE() optimizer = optim.SGD(filter(lambda p: p.requires_grad, model.parameters()), lr=0) from datetime import datetime current_time = datetime.now().strftime('%b%d_%H-%M-%S') trial = Trial( model, optimizer, nn.MSELoss(reduction='sum'), ['acc', 'loss'], pass_state=True, callbacks=[ callbacks.TensorBoardImages(comment=current_time, name='Prediction', write_each_epoch=True, key=torchbearer.Y_PRED, pad_value=1, nrow=16), callbacks.TensorBoardImages(comment=current_time + '_cifar_vae', name='Target', write_each_epoch=False, key=torchbearer.Y_TRUE, pad_value=1, nrow=16) ]).load_state_dict( torch.load(os.path.join(base_dir, file)), resume=False).with_generators( train_generator=testloader, val_generator=testloader).for_train_steps(1).to(device) trial.run( ) # Evaluate doesn't work with tensorboard in torchbearer, seems to have been fixed in most recent version
def run(count, glimpse_size, memory_size, iteration, device='cuda'): base_dir = os.path.join('celeba_' + str(memory_size), str(glimpse_size)) if not os.path.exists(base_dir): os.makedirs(base_dir) transform_train = transforms.Compose([ transforms.ToTensor() ]) dataset = torchvision.datasets.ImageFolder(root='./cropped_celeba/', transform=transform_train) splitter = DatasetValidationSplitter(len(dataset), 0.05) trainset = splitter.get_train_dataset(dataset) # Save the ids torch.save((splitter.train_ids, splitter.valid_ids), os.path.join(base_dir, 'split.dat')) trainloader = torch.utils.data.DataLoader(trainset, batch_size=128, shuffle=True, num_workers=10) model = CelebDraw(count, glimpse_size, memory_size) optimizer = optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=1e-4) from datetime import datetime current_time = datetime.now().strftime('%b%d_%H-%M-%S') call_a = callbacks.TensorBoardImages(comment=current_time, name='Prediction', write_each_epoch=True, key=torchbearer.Y_PRED) call_a.on_step_training = call_a.on_step_validation # Hack to make this log training samples call_b = callbacks.TensorBoardImages(comment=current_time + '_celeba', name='Target', write_each_epoch=True, key=torchbearer.Y_TRUE) call_b.on_step_training = call_b.on_step_validation # Hack to make this log training samples trial = Trial(model, optimizer, nn.MSELoss(reduction='sum'), ['acc', 'loss'], pass_state=True, callbacks=[ joint_kl_divergence(MU, LOGVAR), callbacks.MostRecent(os.path.join(base_dir, 'iter_' + str(iteration) + '.{epoch:02d}.pt')), callbacks.GradientClipping(5), call_a, call_b ]).with_generators(train_generator=trainloader).to(device) trial.run(100)
def draw(count, glimpse_size, memory_size, file, device='cuda'): base_dir = os.path.join('celeba_' + str(memory_size), str(glimpse_size)) transform = transforms.Compose([ transforms.ToTensor() ]) dataset = torchvision.datasets.ImageFolder(root='./cropped_celeba/', transform=transform) splitter = DatasetValidationSplitter(len(dataset), 0.05) # load the ids splitter.train_ids, splitter.valid_ids = torch.load(os.path.join(base_dir, 'split.dat')) testset = splitter.get_val_dataset(dataset) testloader = torch.utils.data.DataLoader(testset, batch_size=128, shuffle=True, num_workers=10) model = CelebDraw(count, glimpse_size, memory_size, output_stages=True) optimizer = optim.SGD(filter(lambda p: p.requires_grad, model.parameters()), lr=0) from datetime import datetime current_time = datetime.now().strftime('%b%d_%H-%M-%S') from visualise import StagesGrid trial = Trial(model, optimizer, nn.MSELoss(reduction='sum'), ['loss'], pass_state=True, callbacks=[ callbacks.TensorBoardImages(comment=current_time, nrow=10, num_images=20, name='Prediction', write_each_epoch=True, key=torchbearer.Y_PRED, pad_value=1), callbacks.TensorBoardImages(comment=current_time + '_celeb', nrow=10, num_images=20, name='Target', write_each_epoch=False, key=torchbearer.Y_TRUE, pad_value=1), callbacks.TensorBoardImages(comment=current_time + '_celeb_mask', nrow=10, num_images=20, name='Masked Target', write_each_epoch=False, key=MASKED_TARGET, pad_value=1), StagesGrid('celeb_stages.png', STAGES, 20) ]).load_state_dict(torch.load(os.path.join(base_dir, file)), resume=False).with_generators(train_generator=testloader, val_generator=testloader).for_train_steps(1).for_val_steps(1).to(device) trial.run() # Evaluate doesn't work with tensorboard in torchbearer, seems to have been fixed in most recent version
def run(iteration, device='cuda:1'): transform_train = transforms.Compose([ transforms.RandomHorizontalFlip(), transforms.ColorJitter(0.25, 0.25, 0.25, 0.25), transforms.ToTensor() ]) transform_test = transforms.Compose([transforms.ToTensor()]) trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform_train) trainloader = torch.utils.data.DataLoader(trainset, batch_size=128, shuffle=True, num_workers=10) testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform_test) testloader = torch.utils.data.DataLoader(testset, batch_size=100, shuffle=False, num_workers=10) base_dir = 'cifar_vae' model = CifarVAE() optimizer = optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=5e-4) from datetime import datetime current_time = datetime.now().strftime('%b%d_%H-%M-%S') trial = Trial( model, optimizer, nn.MSELoss(reduction='sum'), ['acc', 'loss'], pass_state=True, callbacks=[ tm.kl_divergence(MU, LOGVAR, beta=2), callbacks.MultiStepLR([50, 90]), callbacks.MostRecent( os.path.join(base_dir, 'iter_' + str(iteration) + '.{epoch:02d}.pt')), callbacks.GradientClipping(5), callbacks.TensorBoardImages(comment=current_time, name='Prediction', write_each_epoch=True, key=torchbearer.Y_PRED), callbacks.TensorBoardImages(comment=current_time + '_cifar_vae', name='Target', write_each_epoch=False, key=torchbearer.Y_TRUE), ]).with_generators( train_generator=trainloader, val_generator=testloader).for_val_steps(5).to(device) trial.run(100)