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)
