def compute_likelihoods(args, model, dset, weights_name):
if dset == "mnist":
dataset = sample_mnist(args.img_size)
elif dset == "kmnist":
dataset = sample_kmnist(args.img_size)
elif dset == "fmnist":
dataset = sample_fmnist(args.img_size)
elif dset == "notmnist":
dataset = sample_data_one_channel("notMNIST_large/", args.img_size)
n_bins = 2.**args.n_bits
log_likelihoods = []
for i, datapoint in enumerate(dataset):
image, _ = datapoint
image = image.to(device)
log_p, logdet, _ = model(image + torch.rand_like(image) / n_bins)
logdet = logdet.mean()
loss, log_p, log_det = likelihood_loss(log_p, logdet, args.img_size,
n_bins)
log_likelihoods.append(log_p.item() + log_det.item())
print('#{} Loss: {}; logP: {}; logdet: {}'.format(
i, loss.item(), log_p.item(), log_det.item()))
if i % 20 == 0:
np.save(dset + "_" + weights_name + "_likelihoods",
np.array(log_likelihoods))
np.save(dset + "_" + weights_name + "_likelihoods",
np.array(log_likelihoods))
def main():
utils.create_exp_dir(args.save, scripts_to_save=glob.glob('*.py'))
print(args)
# Basic Setup
np.random.seed(args.seed)
torch.manual_seed(args.seed)
torch.cuda.manual_seed(args.seed)
torch.cuda.set_device(2)
cudnn.benchmark = True
cudnn.enabled = True
n_channels = 3
n_bins = 2.**args.n_bits
approx_samples = 4
# Define model
model_single = Network(n_channels,
args.n_flow,
args.n_block,
n_bins,
affine=args.affine,
conv_lu=not args.no_lu)
model = nn.DataParallel(model_single, device_ids=[2, 3])
model = model.to(device)
optimizer = torch.optim.Adam(model.parameters(), args.learning_rate)
dataset = iter(sample_cifar10(args.batch, args.img_size))
# Sample generated images
z_sample = []
z_shapes = calc_z_shapes(n_channels, args.img_size, args.n_flow,
args.n_block)
for z in z_shapes:
z_new = torch.randn(args.n_sample, *z) * args.temp
z_sample.append(z_new.to(device))
with tqdm(range(args.iter)) as pbar:
for i in pbar:
# Training procedure
model.train()
# Get a random minibatch from the search queue with replacement
input, _ = next(dataset)
input = Variable(input,
requires_grad=False).cuda(non_blocking=True)
input = input.repeat(approx_samples, 1, 1, 1)
log_p, logdet, _ = model(input + torch.rand_like(input) / n_bins)
loss, _, _ = likelihood_loss(log_p, logdet, args.img_size, n_bins)
loss_variance = likelihood_loss_variance(log_p, logdet,
args.img_size, n_bins,
approx_samples)
loss = loss + loss_variance
# Optimize model
optimizer.zero_grad()
loss.backward()
optimizer.step()
pbar.set_description("Loss: {}".format(loss.item()))
# Save generated samples
if i % 100 == 0:
with torch.no_grad():
tvutils.save_image(
model_single.reverse(z_sample).cpu().data,
"{}/samples/{}.png".format(args.save,
str(i + 1).zfill(6)),
normalize=False,
nrow=10,
)
# Save checkpoint
if i % 1000 == 0:
model_single.genotype()
torch.save(
model.state_dict(),
"{}/checkpoint/model_{}.pt".format(args.save,
str(i + 1).zfill(6)))
# Save latest weights
utils.save(model, os.path.join(args.save, 'latest_weights.pt'))
def main():
utils.create_exp_dir(args.save, scripts_to_save=glob.glob('*.py'))
print(args)
seed = random.randint(1, 100000000)
print(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
torch.cuda.set_device(args.gpu)
cudnn.benchmark = True
cudnn.enabled = True
n_channels = 3
n_bins = 2.**args.n_bits
# Define model and loss criteria
model = SearchNetwork(n_channels,
args.n_flow,
args.n_block,
n_bins,
affine=args.affine,
conv_lu=not args.no_lu)
model = nn.DataParallel(model, [args.gpu])
model.load_state_dict(
torch.load("architecture.pt", map_location="cuda:{}".format(args.gpu)))
model = model.module
genotype = model.sample_architecture()
with open(args.save + '/genotype.pkl', 'wb') as fp:
pickle.dump(genotype, fp)
model_single = EnsembleNetwork(n_channels,
args.n_flow,
args.n_block,
n_bins,
genotype,
affine=args.affine,
conv_lu=not args.no_lu)
model = model_single
model = model.to(device)
optimizer = torch.optim.Adam(model.parameters(), args.learning_rate)
dataset = iter(sample_cifar10(args.batch, args.img_size))
# Sample generated images
z_sample = []
z_shapes = calc_z_shapes(n_channels, args.img_size, args.n_flow,
args.n_block)
for z in z_shapes:
z_new = torch.randn(args.n_sample, *z) * args.temp
z_sample.append(z_new.to(device))
with tqdm(range(args.iter)) as pbar:
for i in pbar:
# Training procedure
model.train()
# Get a random minibatch from the search queue with replacement
input, _ = next(dataset)
input = Variable(input,
requires_grad=False).cuda(non_blocking=True)
log_p, logdet, _ = model(input + torch.rand_like(input) / n_bins)
logdet = logdet.mean()
loss, _, _ = likelihood_loss(log_p, logdet, args.img_size, n_bins)
# Optimize model
optimizer.zero_grad()
loss.backward()
optimizer.step()
pbar.set_description("Loss: {}".format(loss.item()))
# Save generated samples
if i % 100 == 0:
with torch.no_grad():
tvutils.save_image(
model_single.reverse(z_sample).cpu().data,
"{}/samples/{}.png".format(args.save,
str(i + 1).zfill(6)),
normalize=False,
nrow=10,
)
# Save checkpoint
if i % 1000 == 0:
utils.save(model, os.path.join(args.save, 'latest_weights.pt'))