def main():
device = torch.device(ARGS.device)
if ARGS.t:
data = bmnist(batch_size=1)[:2] # ignore test split
else:
data = bmnist()[:2] # ignore test split
model = VAE(z_dim=ARGS.zdim, deep=True, device=device)
model.to(device)
optimizer = torch.optim.Adam(model.parameters())
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, ARGS.epochs+20)
save_samples(model, 10, "Samples from VAE, {}-D latent space, epoch = {}".format(ARGS.zdim, 0),
"./results/test_epoch_0.pdf")
train_curve, val_curve = [], []
for epoch in range(ARGS.epochs):
elbos = run_epoch(model, data, optimizer, device)
train_elbo, val_elbo = elbos
train_curve.append(train_elbo)
val_curve.append(val_elbo)
print("[Epoch {}] train elbo: {}, val_elbo: {}, learning rate: {}".format(
epoch, train_elbo, val_elbo, optimizer.state_dict()['param_groups'][0]['lr']))
scheduler.step()
# --------------------------------------------------------------------
# Add functionality to plot samples from model during training.
# You can use the make_grid functioanlity that is already imported.
# --------------------------------------------------------------------
# Biggest change in sample quality is first couple of steps
if (epoch < 5) or ((epoch+1) % 10) == 0:
if ARGS.t:
save_samples(model, 10, "Samples from VAE, {}-D latent space, epoch = {}".format(ARGS.zdim, epoch+1),
"./results/test_epoch_{}.pdf".format(epoch+1))
else:
save_samples(model, 10, "Samples from VAE, {}-D latent space, epoch = {}".format(ARGS.zdim, epoch+1),
"./results/samples_epoch_{}.pdf".format(epoch+1))
# --------------------------------------------------------------------
# Add functionality to plot plot the learned data manifold after
# if required (i.e., if zdim == 2). You can use the make_grid
# functionality that is already imported.
# --------------------------------------------------------------------
if ARGS.zdim == 2:
GRID_SIZE = 20
means = model.get_latent(GRID_SIZE)
grid = make_grid(means.view(GRID_SIZE**2, 1, 28, -1).permute(0, 1, 3, 2), nrow = GRID_SIZE)
plt.figure(figsize=(12,12))
plt.imshow(grid.permute(2, 1, 0).numpy())
plt.title("Learned MNIST manifold")
plt.savefig("./results/mnist_manifold.pdf".format(epoch+1))
if not ARGS.t:
save_elbo_plot(train_curve, val_curve, 'elbo.pdf')
def main():
writer = VaeWriter(ARGS.output_dir)
device = 'cuda' if torch.cuda.is_available() else 'cpu'
writer.log("Device: " + device)
data = bmnist()[:2] # ignore test split
model = VAE(z_dim=ARGS.zdim).to(device)
optimizer = torch.optim.Adam(model.parameters())
writer.log("Model:\n" + str(model))
for epoch in range(ARGS.epochs):
elbos = run_epoch(model, data, optimizer)
train_elbo, val_elbo = elbos
# Output stats and images
writer.save_stats(train_elbo, val_elbo)
writer.log(
f"[Epoch {epoch}] train elbo: {train_elbo} val_elbo: {val_elbo}")
writer.save_elbo_plot()
sample_imgs, _ = model.sample(25)
writer.save_images(sample_imgs, epoch)
# Print manifold
if ARGS.zdim == 2:
side_len = 20
eps = 1e-9
zs = np.linspace(0 + eps, 1 - eps, side_len)
samples = torch.stack(
[torch.Tensor(ndtri([x, y])) for x in zs for y in zs])
manifold = model.decoder(samples.to(device)).view(
-1, 1, side_len, side_len)
writer.save_manifold(manifold, epoch)
def main(epochs, zdim, _run):
# Take train and val sets
data = bmnist()[:2]
model = VAE(z_dim=zdim).to(device)
optimizer = torch.optim.Adam(model.parameters())
for epoch in range(1, epochs + 1):
# Save samples at beginning, 50% and 100% of training
if int(100 * epoch / epochs) in [int(100 / epochs), 50, 100]:
fname = 'vae_{:d}.png'.format(epoch)
save_samples(model, fname)
losses = run_epoch(model, data, optimizer)
train_rec, train_kl, val_rec, val_kl = losses
train_elbo = train_rec + train_kl
val_elbo = val_rec + val_kl
print(f"[Epoch {epoch}] train elbo: {train_elbo} val_elbo: {val_elbo}")
_run.log_scalar('train_elbo', train_elbo, epoch)
_run.log_scalar('val_elbo', val_elbo, epoch)
_run.log_scalar('train_rec', train_rec, epoch)
_run.log_scalar('train_kl', train_kl, epoch)
_run.log_scalar('val_rec', val_rec, epoch)
_run.log_scalar('val_kl', val_kl, epoch)
if zdim == 2:
save_manifold(model)
def main(args):
flags = ('epochs', 'zdim')
(epochs, zdim) = itemgetter(*flags)(vars(args))
data = bmnist()[:2] # ignore test split
model = VAE(item_dims, z_dim=zdim)
optimizer = torch.optim.Adam(model.parameters())
path = './vae'
mkdir(path)
plot(model, zdim, path, epoch=0)
train_curve, val_curve = [], []
for epoch in trange(epochs):
elbos = run_epoch(model, data, optimizer)
train_elbo, val_elbo = elbos
train_curve.append(train_elbo)
val_curve.append(val_elbo)
print(f'[Epoch {epoch}] train elbo: {train_elbo} val_elbo: {val_elbo}')
# --------------------------------------------------------------------
# Add functionality to plot samples from model during training.
# You can use the make_grid functionality that is already imported.
# --------------------------------------------------------------------
plot(model, zdim, path, epoch + 1)
# --------------------------------------------------------------------
# Add functionality to plot the learned data manifold after
# if required (i.e., if zdim == 2). You can use the make_grid
# functionality that is already imported.
# -------------------------------------------------------------------
save_elbo_plot(train_curve, val_curve, 'elbo.png')
def main():
data = bmnist()[:2] # ignore test split
model = VAE(z_dim=ARGS.zdim)
optimizer = torch.optim.Adam(model.parameters())
size_width = 28
train_curve, val_curve = [], []
for epoch in range(ARGS.epochs):
elbos = run_epoch(model, data, optimizer)
train_elbo, val_elbo = elbos
train_curve.append(train_elbo)
val_curve.append(val_elbo)
print(f"[Epoch {epoch}] train elbo: {train_elbo} val_elbo: {val_elbo}")
# --------------------------------------------------------------------
# Add functionality to plot samples from model during training.
# You can use the make_grid functioanlity that is already imported.
# --------------------------------------------------------------------
_, mean_sample = model.sample(64)
save_sample(mean_sample, size_width, epoch)
if ARGS.zdim == 2:
print("manifold")
manifold = model.manifold_sample(256)
save_sample(manifold, size_width, epoch, 16)
np.save('curves.npy', {'train': train_curve, 'val': val_curve})
# --------------------------------------------------------------------
# Add functionality to plot plot the learned data manifold after
# if required (i.e., if zdim == 2). You can use the make_grid
# functionality that is already imported.
# --------------------------------------------------------------------
save_elbo_plot(train_curve, val_curve, 'elbo.pdf')
def main():
data = bmnist()[:2] # ignore test split
model = VAE(z_dim=ARGS.zdim)
model=model.cuda()
optimizer = torch.optim.Adam(model.parameters())
train_curve, val_curve = [], []
for epoch in range(ARGS.epochs):
elbos = run_epoch(model, data, optimizer)
train_elbo, val_elbo = elbos
train_curve.append(train_elbo)
val_curve.append(val_elbo)
print(f"[Epoch {epoch}] train elbo: {train_elbo} val_elbo: {val_elbo}")
# --------------------------------------------------------------------
# Add functionality to plot samples from model during training.
# You can use the make_grid functioanlity that is already imported.
# --------------------------------------------------------------------
# --------------------------------------------------------------------
# Add functionality to plot plot the learned data manifold after
# if required (i.e., if zdim == 2). You can use the make_grid
# functionality that is already imported.
# --------------------------------------------------------------------
##save_elbo_plot(train_curve, val_curve, 'elbo.pdf')
model.sample(100)
def main():
# general setup
os.makedirs('./figures', exist_ok=True)
imw = 28
device = torch.device(ARGS.device)
data = bmnist(root=ARGS.data, batch_size=ARGS.batch_size, download=False)[:2] # ignore test split
model = VAE(x_dim=imw**2, z_dim=ARGS.zdim, device=device)
optimizer = torch.optim.Adam(model.parameters())
train_curve, val_curve = [], []
for epoch in range(ARGS.epochs):
(train_elbo, val_elbo) = run_epoch(model, data, optimizer, device=device)
train_curve.append(train_elbo)
val_curve.append(val_elbo)
print(f"[Epoch {epoch}] train elbo: {train_elbo:.4e} val_elbo: {val_elbo:.4e}")
_, μ_sample = model.sample(ARGS.samples)
save_sample(μ_sample, imw, epoch)
if ARGS.zdim == 2:
manifold = model.manifold(ARGS.manifold_samples)
save_sample(manifold, imw, epoch, ARGS.manifold_samples, 'manifold')
torch.save({'train': train_curve, 'val': val_curve}, f"vae_{ARGS.zdim}_curves.pt")
def main():
data = bmnist()[:2] # ignore test split
device = 'cuda' if torch.cuda.is_available() else 'cpu'
model = VAE(z_dim=ARGS.zdim, device=device)
optimizer = torch.optim.Adam(model.parameters())
# sample before training
sample_model(model, 0, 9999, 9999)
train_curve, val_curve = [], []
for epoch in range(ARGS.epochs):
elbos = run_epoch(model, data, optimizer)
train_elbo, val_elbo = elbos
train_curve.append(train_elbo)
val_curve.append(val_elbo)
print(f"[Epoch {epoch}] train elbo: {train_elbo} val_elbo: {val_elbo}")
# --------------------------------------------------------------------
# Add functionality to plot samples from model during training.
# You can use the make_grid functioanlity that is already imported.
# --------------------------------------------------------------------
sample_model(model, epoch, train_elbo, val_elbo)
# --------------------------------------------------------------------
# Add functionality to plot plot the learned data manifold after
# if required (i.e., if zdim == 2). You can use the make_grid
# functionality that is already imported.
# --------------------------------------------------------------------
if ARGS.zdim == 2:
plot_manifold(model)
save_elbo_plot(train_curve, val_curve, 'elbo.pdf')
def main():
data = bmnist()[:2] # ignore test split
model = VAE(z_dim=ARGS.zdim)
optimizer = torch.optim.Adam(model.parameters())
train_curve, val_curve = [], []
for epoch in range(ARGS.epochs):
if epoch == 0:
samples = model.sample(9)[0]
save_image(samples.view(9, 1, 28, 28),
"output_vae_epoch_{}.png".format(epoch),
nrow=3)
elbos = run_epoch(model, data, optimizer)
train_elbo, val_elbo = elbos
train_curve.append(train_elbo)
val_curve.append(val_elbo)
print(f"[Epoch {epoch}] train elbo: {train_elbo} val_elbo: {val_elbo}")
if epoch == ARGS.epochs / 2 or epoch == ARGS.epochs - 1:
samples = model.sample(9)[0]
save_image(samples.view(9, 1, 28, 28),
"output_vae_epoch_{}.png".format(epoch),
nrow=3)
# --------------------------------------------------------------------
# Add functionality to plot plot the learned data manifold after
# if required (i.e., if zdim == 2). You can use the make_grid
# functionality that is already imported.
# --------------------------------------------------------------------
save_elbo_plot(train_curve, val_curve, 'elbo.pdf')
def main():
os.makedirs('images_vae', exist_ok=True)
data = bmnist()[:2] # ignore test split
model = VAE(z_dim=ARGS.zdim)
optimizer = torch.optim.Adam(model.parameters())
train_curve, val_curve = [], []
for epoch in range(ARGS.epochs):
elbos = run_epoch(model, data, optimizer, epoch)
train_elbo, val_elbo = elbos
train_curve.append(train_elbo)
val_curve.append(val_elbo)
print(f"[Epoch {epoch}] train elbo: {train_elbo} val_elbo: {val_elbo}")
# --------------------------------------------------------------------
# Add functionality to plot samples from model during training.
# You can use the make_grid functioanlity that is already imported.
# --------------------------------------------------------------------
samples, samples_mean = model.sample(25)
save_imgs = samples.view(-1, 1, 28, 28).cpu()
save_image(save_imgs.data[:25],
'images_vae/%d.png' % epoch,
nrow=5,
normalize=True)
# --------------------------------------------------------------------
# Add functionality to plot plot the learned data manifold after
# if required (i.e., if zdim == 2). You can use the make_grid
# functionality that is already imported.
# --------------------------------------------------------------------
#print(samples_mean)
save_elbo_plot(train_curve, val_curve, 'elbo.pdf')
def main():
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
# Print all configs to confirm parameter settings
print_flags()
data = bmnist()[:2] # ignore test split
model = VAE(z_dim=ARGS.zdim, device=device).to(device)
optimizer = torch.optim.Adam(model.parameters())
# Create output directories
os.makedirs('./images/vae', exist_ok=True)
train_curve, val_curve = [], []
for epoch in range(ARGS.epochs):
elbos = run_epoch(model, data, optimizer, device)
train_elbo, val_elbo = elbos
train_curve.append(train_elbo)
val_curve.append(val_elbo)
print(f"[Epoch {epoch}] train elbo: {train_elbo} val_elbo: {val_elbo}")
# --------------------------------------------------------------------
# Add functionality to plot samples from model during training.
# You can use the make_grid functioanlity that is already imported.
# --------------------------------------------------------------------
if ARGS.zdim != 2:
im_samples, means_samples = model.sample(25)
save_image(im_samples.view(im_samples.shape[0], 1, 28, 28),
'./images/vae/bernoulli_{}.png'.format(epoch),
nrow=5,
normalize=True)
save_image(means_samples.view(means_samples.shape[0], 1, 28, 28),
'./images/vae/mean_{}.png'.format(epoch),
nrow=5,
normalize=True)
# --------------------------------------------------------------------
# Add functionality to plot plot the learned data manifold after
# if required (i.e., if zdim == 2). You can use the make_grid
# functionality that is already imported.
# --------------------------------------------------------------------
if ARGS.zdim == 2:
model.eval()
x_values = norm.ppf(np.linspace(0.0001, 0.9999, 15))
y_values = norm.ppf(np.linspace(0.0001, 0.9999, 15))
manifold = torch.FloatTensor(
np.array(np.meshgrid(x_values, y_values)).T).view(-1, 2).to(device)
with torch.no_grad():
imgs = model.decoder(manifold)
save_image(imgs.view(imgs.shape[0], 1, 28, 28),
'./images/vae/manifold.png'.format(epoch),
nrow=15,
normalize=True)
save_elbo_plot(train_curve, val_curve,
'./images/vae/elbo_zdim-{}.png'.format(ARGS.zdim))
save_elbo_plot(train_curve, val_curve,
'./images/vae/elbo_zdim-{}.eps'.format(ARGS.zdim))
def main():
data = bmnist()[:2] # ignore test split
model = VAE(z_dim=ARGS.zdim)
optimizer = torch.optim.Adam(model.parameters())
train_curve, val_curve = [], []
for epoch in range(ARGS.epochs):
elbos = run_epoch(model, data, optimizer)
train_elbo, val_elbo = elbos
train_curve.append(train_elbo)
val_curve.append(val_elbo)
print(f'[{datetime.now().strftime("%Y-%m-%d %H:%M")}] Epoch {epoch:02d}, train elbo: {train_elbo:07f}, val_elbo: {val_elbo:07f}')
# --------------------------------------------------------------------
# Add functionality to plot samples from model during training.
# You can use the make_grid functioanlity that is already imported.
# --------------------------------------------------------------------
# plot the samples
sampled_ims, im_means = model.sample(16)
save_samples_plot(sampled_ims, im_means, model, epoch, ARGS.zdim)
# save the results
results = {
'train elbo': train_curve,
'val elbo': val_curve
}
results_df = df.from_dict(results)
results_path = RESULTS_DIR / f'{model.__class__.__name__}_results.csv'
results_df.to_csv(results_path, sep=';', encoding='utf-8', index=False)
# save the model
model_path = RESULTS_DIR / f'{model.__class__.__name__}_model.pt'
torch.save(model.state_dict(), model_path)
# --------------------------------------------------------------------
# Add functionality to plot plot the learned data manifold after
# if required (i.e., if zdim == 2). You can use the make_grid
# functionality that is already imported.
# --------------------------------------------------------------------
if ARGS.zdim == 2:
# create data manifold
n_rows = 20
x = y = np.linspace(0, 1, n_rows)
samples = [torch.erfinv(2 * torch.tensor([x, y], device='cpu') - 1) * np.sqrt(2) for x in x for y in y]
samples = torch.stack(samples)
manifold = model.decoder(samples).view(-1, 1, 28, 28)
grid = make_grid(manifold, nrow=n_rows)
plt.imsave(RESULTS_DIR / 'VAE_manifold.png', grid.cpu().numpy().transpose(1, 2, 0))
plot_path = RESULTS_DIR / f'{model.__class__.__name__}_elbo.png'
save_elbo_plot(train_curve, val_curve, plot_path)
def main():
if torch.cuda.is_available():
device = 'cuda'
else:
device = 'cpu'
data = bmnist()[:2] # ignore test split
model = VAE(z_dim=ARGS.zdim, device=device).to(device)
optimizer = torch.optim.Adam(model.parameters())
train_curve, val_curve = [], []
for epoch in range(ARGS.epochs):
elbos = run_epoch(model, data, optimizer, device)
train_elbo, val_elbo = elbos
train_curve.append(train_elbo)
val_curve.append(val_elbo)
print(f"[Epoch {epoch}] train elbo: {train_elbo} val_elbo: {val_elbo}")
if True:
path = 'images/ELBO_VAE.png'
plt.plot(train_curve, label='training ELBO')
plt.plot(val_curve, label='validation ELBO')
plt.title("ELBO VAE")
plt.xlabel("Epochs")
plt.ylabel("ELBO")
plt.legend()
plt.savefig(path)
plt.close('all')
if True:
sampled_ims, im_means = model.sample(100)
# print(sampled_ims.shape, im_means.shape)
# print(asdad)
save_image(sampled_ims,
'images/VAE_epoch_{}.png'.format(epoch),
nrow=10,
normalize=True)
if True:
sampled_ims = model.plot_2d_manifold()
save_image(sampled_ims,
'images/VAE_manifold_epoch_{}.png'.format(epoch),
nrow=10,
normalize=True)
if False:
sampled_ims = model.plot_beta_VAE_manifold()
save_image(sampled_ims,
'images/BETA_VAE_manifold_epoch_{}.png'.format(epoch),
nrow=10,
normalize=True)
save_elbo_plot(train_curve, val_curve, 'elbo.pdf')
def main():
data = bmnist()[:2] # ignore test split
model = VAE(z_dim=ARGS.zdim)
optimizer = torch.optim.Adam(model.parameters())
start = datetime.now().strftime("%Y-%m-%d_%H:%M")
f_train = open('elbos_vae/train_' + start + '.txt', 'w')
f_valid = open('elbos_vae/valid_' + start + '.txt', 'w')
sample_images, _ = model.sample(16)
grid = make_grid(sample_images.detach(), nrow=4, padding=0)
save_image(grid, 'samples_vae/start_' + str(start) + '.png')
train_curve, val_curve = [], []
for epoch in range(ARGS.epochs):
elbos = run_epoch(model, data, optimizer)
train_elbo, val_elbo = elbos
print(f"[Epoch {epoch}] train elbo: {train_elbo} val_elbo: {val_elbo}")
f_train.write(', ' + str(train_elbo))
f_valid.write(', ' + str(val_elbo))
if epoch == ARGS.epochs / 2 - 1 or epoch == ARGS.epochs / 2 - 0.5:
sample_images, _ = model.sample(16)
grid = make_grid(sample_images.detach(), nrow=4, padding=0)
save_image(grid, 'samples_vae/middle_' + str(start) + '.png')
if ARGS.zdim == 2:
with torch.no_grad():
grid = np.empty((420, 420))
for i, y in enumerate(np.linspace(.05, .95, 15)):
for j, x in enumerate(np.linspace(.05, .95, 15)):
# Use ppf to sample the correct spot that matches the 0.05 to 0.95 area
noise = torch.tensor(np.array([[norm.ppf(x), norm.ppf(y)]])\
.astype('float32'))
means = model._decoder(noise)
grid[(14-i)*28:(15-i)*28, j*28:(j+1)*28] = means\
.reshape(28, 28).data.numpy()
plt.figure(figsize=(8, 10))
plt.imshow(grid, origin="upper", cmap="gray")
plt.tight_layout()
plt.savefig('manifold_' + str(start) + '.png')
sample_images, _ = model.sample(16)
grid = make_grid(sample_images.detach(), nrow=4, padding=0)
save_image(grid, 'samples_vae/end_' + str(start) + '.png')
# save_grid(sample_images.detach(), 'samples/end_' + str(start))
f_train.close()
f_valid.close()
def main(device):
data = bmnist()[:2] # ignore test split
model = VAE(z_dim=ARGS.zdim, device=device, input_dim=ARGS.input_dim)
model = model.to(device)
print(model.device)
optimizer = torch.optim.Adam(model.parameters())
# size by size for grid or manifold
grid_size = 4
manifold_size = 20
x_and_y_dim = int(math.sqrt(ARGS.input_dim))
sampled, im_means = model.sample(grid_size * grid_size, device)
plot_samples(sampled, x_and_y_dim, grid_size,
"samples_before_training.png")
#plot_samples(im_means, x_and_y_dim, grid_size, "sample_means_before_training.png")
train_curve, val_curve = [], []
for epoch in range(ARGS.epochs):
elbos = run_epoch(model, data, optimizer)
train_elbo, val_elbo = elbos
train_curve.append(train_elbo)
val_curve.append(val_elbo)
print(f"[Epoch {epoch}] train elbo: {train_elbo} val_elbo: {val_elbo}")
# --------------------------------------------------------------------
# Add functionality to plot samples from model during training.
# You can use the make_grid functioanlity that is already imported.
# --------------------------------------------------------------------
sampled, im_means = model.sample(grid_size * grid_size, device)
file_name = "samples_epoch_" + str(epoch) + ".png"
plot_samples(sampled, x_and_y_dim, grid_size, file_name)
#plot_samples(im_means, x_and_y_dim, grid_size, file_name = "sample_means_" + str(epoch) + ".png")
#plot_manifold(x_and_y_dim, manifold_size, model)
# --------------------------------------------------------------------
# Add functionality to plot plot the learned data manifold after
# if required (i.e., if zdim == 2). You can use the make_grid
# functionality that is already imported.
# --------------------------------------------------------------------
sampled, im_means = model.sample(grid_size * grid_size, device)
plot_samples(sampled, x_and_y_dim, grid_size, "samples_after_training.png")
#plot_samples(im_means, x_and_y_dim, grid_size, "samples_means_after_training.png")
if ARGS.zdim == 2:
plot_manifold(x_and_y_dim, manifold_size, model)
save_elbo_plot(train_curve, val_curve, 'elbo.pdf')
def main():
# ID for the current run
run_id = datetime.now().strftime("vae_%Y-%m-%d_%H-%M-%S")
dir_path = 'vae/' + run_id + '/'
if not os.path.exists(dir_path):
os.makedirs(dir_path)
data = bmnist()[:2] # ignore test split
model = VAE(z_dim=ARGS.zdim).to(device) # send to cuda if available
optimizer = torch.optim.Adam(model.parameters())
train_curve, val_curve = [], []
for epoch in range(ARGS.epochs):
elbos = run_epoch(model, data, optimizer)
train_elbo, val_elbo = elbos
model.train()
train_curve.append(train_elbo)
model.eval()
val_curve.append(val_elbo)
print(
f"[Epoch {epoch+1:02d}] train elbo: {train_elbo:.2f} val_elbo: {val_elbo:.2f}"
)
# plot samples every 5 epochs
if ((epoch + 1) % 5 == 0) or epoch == 0:
sampled_ims, im_means = model.sample(25)
sample_plot = make_grid(im_means,
nrow=5).cpu().detach().numpy().transpose(
1, 2, 0)
path = f'samples_epoch_{epoch:02d}.png'
plt.imsave(dir_path + path, sample_plot)
# plot manifold at the end of the training
if ARGS.zdim == 2:
e = 1e-3
linspace = np.linspace(0 + e, 1 - e, 20)
X = norm.ppf(linspace)
Y = norm.ppf(linspace)
z = torch.tensor([[x, y] for x in X for y in Y])
sampled_ims, im_means = model.sample(20 * 20, z)
sample_plot = make_grid(im_means,
nrow=20).cpu().detach().numpy().transpose(
1, 2, 0)
path = f'manifold.png'
plt.imsave(dir_path + path, sample_plot)
save_elbo_plot(train_curve, val_curve, dir_path + 'elbo.pdf')
def main():
data = bmnist()[:2] # ignore test split
model = VAE(z_dim=ARGS.zdim)
optimizer = torch.optim.Adam(model.parameters())
# store images in result folder
result_folder = './results'
train_curve, val_curve = [], []
for epoch in range(ARGS.epochs):
elbos = run_epoch(model, data, optimizer)
train_elbo, val_elbo = elbos
train_curve.append(train_elbo)
val_curve.append(val_elbo)
print(f"[Epoch {epoch}] train elbo: {train_elbo} val_elbo: {val_elbo}")
# --------------------------------------------------------------------
# Add functionality to plot samples from model during training.
# You can use the make_grid functioanlity that is already imported.
# --------------------------------------------------------------------
n_samples = 25 # allows for 5x5 grid of samples
# get samples from model
sampled_ims, im_means = model.sample(n_samples)
# reshape samples, means
sampled_ims = sampled_ims.view(n_samples, 1, 28, 28)
im_means = im_means.view(n_samples, 1, 28, 28)
# save sampled images
save_image(sampled_ims,
result_folder + '/samples/' +
'samples_iteration_{}.png'.format(epoch),
nrow=5,
normalize=True,
padding=5)
# save sampled means
save_image(im_means,
result_folder + '/means/' +
'means_iteration_{}.png'.format(epoch),
nrow=5,
normalize=True,
padding=5)
# --------------------------------------------------------------------
# Add functionality to plot plot the learned data manifold after
# if required (i.e., if zdim == 2). You can use the make_grid
# functionality that is already imported.
# --------------------------------------------------------------------
save_elbo_plot(train_curve, val_curve, 'elbo.pdf')
def main(config):
# create directory for output files (images, graphs and model)
if not os.path.exists(ARGS.output_dir):
os.makedirs(ARGS.output_dir)
# check if GPU is available
device = 'cuda' if torch.cuda.is_available() else 'cpu'
print(f"RUNNING ON {device}")
# load data and configure model
data = bmnist()[:2] # ignore test split
model = VAE(z_dim=config.zdim, device=device)
model.to(device)
optimizer = torch.optim.Adam(model.parameters())
train_curve, val_curve = [], []
for epoch in range(config.epochs):
elbos = run_epoch(model, data, optimizer)
train_elbo, val_elbo = elbos
train_curve.append(train_elbo)
val_curve.append(val_elbo)
print(f"[Epoch {epoch}] train elbo: {train_elbo} val_elbo: {val_elbo}")
# --------------------------------------------------------------------
# Add functionality to plot samples from model during training.
# You can use the make_grid functioanlity that is already imported.
# --------------------------------------------------------------------
with torch.no_grad():
samples, im_means = model.sample(9)
samples = make_grid(samples, nrow=3)
save(
samples,
path=
f"{ARGS.output_dir}/samples_epoch_{epoch}_train_{train_elbo:.2f}_val_{val_elbo:.2f}.jpg"
)
torch.save(model, f'{ARGS.output_dir}/model_epoch_{ARGS.epochs}.pt')
# --------------------------------------------------------------------
# Add functionality to plot plot the learned data manifold after
# if required (i.e., if zdim == 2). You can use the make_grid
# functionality that is already imported.
#
# NB: I implemented this in a seperate script, see manifold.py
# --------------------------------------------------------------------
save_elbo_plot(train_curve, val_curve, f'{ARGS.output_dir}/elbo.pdf')
def main():
"""
Runs all code
:return:
"""
data = bmnist()[:2] # ignore test split
model = VAE(z_dim=ARGS.zdim)
optimizer = torch.optim.Adam(model.parameters())
train_curve, val_curve = [], []
# Plot before training -------------------------------
binary_samples, mean_samples = model.sample(25)
save_samples(binary_samples,
train_elbo=0,
val_elbo=0,
epoch=0,
binaryormean="binary")
save_samples(mean_samples,
train_elbo=0,
val_elbo=0,
epoch=0,
binaryormean="binary")
# ----------------------------------------------------
for epoch in range(ARGS.epochs):
elbos = run_epoch(model, data, optimizer)
train_elbo, val_elbo = elbos
train_curve.append(train_elbo)
val_curve.append(val_elbo)
print(f"[Epoch {epoch}] train elbo: {train_elbo} val_elbo: {val_elbo}")
# save samples during training ---------------------------------------
binary_samples, mean_samples = model.sample(25)
save_samples(binary_samples, train_elbo, val_elbo, epoch, "binary")
save_samples(mean_samples, train_elbo, val_elbo, epoch, "mean")
# --------------------------------------------------------------------
# Plot manifold if number of laten dimensions is 2 -------------------
if ARGS.zdim == 2:
manifold_samples = model.sample_manifold(25, 0.05, 0.95)
save_samples(manifold_samples,
train_elbo=0,
val_elbo=0,
epoch=epoch,
binaryormean="manifold")
# --------------------------------------------------------------------
save_elbo_plot(train_curve, val_curve, 'elbo.pdf')
def main():
data = bmnist()[:2] # ignore test split
model = VAE(z_dim=ARGS.zdim).to(DEVICE)
optimizer = torch.optim.Adam(model.parameters())
os.makedirs('images_vae', exist_ok=True)
train_curve, val_curve = [], []
for epoch in range(ARGS.epochs):
elbos = run_epoch(model, data, optimizer)
train_elbo, val_elbo = elbos
train_curve.append(train_elbo)
val_curve.append(val_elbo)
print(f"[Epoch {epoch}] train elbo: {train_elbo} val_elbo: {val_elbo}")
# --------------------------------------------------------------------
# Add functionality to plot samples from model during training.
# You can use the make_grid functioanlity that is already imported.
# --------------------------------------------------------------------
ims_per_row = 5
sampled_ims, _ = model.sample(ims_per_row * ims_per_row)
grid = make_grid(sampled_ims, nrow=ims_per_row)
save_image(grid,
'images_vae/epoch{}_{}z.png'.format(epoch, ARGS.zdim),
normalize=True)
torch.save(model.state_dict(),
"models/VAE_{}epochs_{}z.pt".format(ARGS.epochs, ARGS.zdim))
# --------------------------------------------------------------------
# Add functionality to plot plot the learned data manifold after
# if required (i.e., if zdim == 2). You can use the make_grid
# functionality that is already imported.
# --------------------------------------------------------------------
if ARGS.zdim == 2:
with torch.no_grad():
steps = 20
density_points = torch.linspace(0, 1, steps)
# Basically use adaptation of torch.distributions.icdf here for manifold z's
z_tensors = [
torch.erfinv(2 * torch.tensor([x, y]) - 1) * np.sqrt(2)
for x in density_points for y in density_points
]
z = torch.stack(z_tensors).to(DEVICE)
_, manifold = model.sample(1, z)
image = make_grid(manifold, nrow=steps)
save_image(image, "images_vae/manifold.pdf")
save_elbo_plot(train_curve, val_curve, 'elbo.pdf')
def main():
if not os.path.exists("samples_vae"):
os.makedirs("samples_vae")
if not os.path.exists("manifolds_vae"):
os.makedirs("manifolds_vae")
data = bmnist(batch_size=ARGS.batch_size)[:2] # ignore test split
model = VAE(z_dim=ARGS.zdim)
model.to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=ARGS.lr_rate)
# Pre-training samples
samples, samples_mu = model.sample(64)
grid = make_grid(samples.unsqueeze(1))
save_image(grid, f"samples_vae/0.png")
train_curve, val_curve = [], []
for epoch in range(ARGS.epochs):
elbos = run_epoch(model, data, optimizer)
train_elbo, val_elbo = elbos
train_curve.append(train_elbo)
val_curve.append(val_elbo)
print(f"[Epoch {epoch}] train elbo: {train_elbo} val_elbo: {val_elbo}")
# --------------------------------------------------------------------
# Add functionality to plot samples from model during training.
# You can use the make_grid functioanlity that is already imported.
# --------------------------------------------------------------------
# We are still in model.eval()
samples, samples_mu = model.sample(64)
grid = make_grid(samples.unsqueeze(1))
save_image(grid, f"samples_vae/{epoch+1}.png")
if ARGS.zdim == 2:
samples, samples_mu = model.sample_manifold(20)
grid = make_grid(samples.unsqueeze(1), nrow=20)
save_image(grid, f"manifolds_vae/{epoch+1}.png")
# --------------------------------------------------------------------
# Add functionality to plot plot the learned data manifold after
# if required (i.e., if zdim == 2). You can use the make_grid
# functionality that is already imported.
# --------------------------------------------------------------------
# I do this in every epoch, to see how the manifold evolves with the training
save_elbo_plot(train_curve, val_curve, 'elbo.pdf')
def main():
device = torch.device(ARGS.device)
data = bmnist()[:2] # ignore test split
model = VAE(z_dim=ARGS.zdim, device=device).to(device)
torch.save(model.state_dict(), ARGS.model_file + "0.pt")
if ARGS.use == "train":
#which epochs to save; first and last already included
save_inter_epochs = np.array([5, 10, 20]) #could make ARG
results = open(ARGS.out, "w+")
results.write(
f"#learning_rate : {ARGS.lr}\n#z_dim : {ARGS.zdim}\n#epochs : {ARGS.epochs}\n"
)
results.write("#epoch train_elbo val_elbo\n")
optimizer = torch.optim.Adam(model.parameters(), lr=ARGS.lr)
train_curve, val_curve = [], []
for epoch in range(ARGS.epochs):
if np.sum(epoch == save_inter_epochs):
torch.save(model.state_dict(), ARGS.model_file + f"{epoch}.pt")
elbos = run_epoch(model, data, optimizer, device)
train_elbo, val_elbo = elbos
train_curve.append(train_elbo)
val_curve.append(val_elbo)
print(
f"[Epoch {epoch}] train elbo: {train_elbo} val_elbo: {val_elbo}"
)
results.write(f"{epoch} {train_elbo} {val_elbo}\n")
torch.save(model.state_dict(), ARGS.model_file + f"{ARGS.epochs}.pt")
results.close()
else:
model.load_state_dict(
torch.load(ARGS.load, map_location=lambda storage, loc: storage))
model.eval()
N = 40 #could make ARG
B = 8 #could make ARG
imgs, means = model.sample(int(N / 2.))
plt.figure(figsize=(B * 1.5, (int(N / B) + 1) * 1.5))
for i in range(int(N / 2.)):
plt.subplot(int(N / B) + 1, B, i * 2 + 1)
plt.imshow(means[i].view(28, 28), cmap="gray")
plt.axis('off')
plt.subplot(int(N / B) + 1, B, i * 2 + 2)
plt.imshow(imgs[i].view(28, 28), cmap="binary")
plt.axis('off')
plt.show()
def main():
data = bmnist()[:2] # ignore test split
# 28 is image dimension
model = VAE(28 * 28, hidden_dim=ARGS.h_dim, z_dim=ARGS.zdim)
optimizer = torch.optim.Adam(model.parameters())
train_curve, val_curve = [], []
for epoch in range(ARGS.epochs):
elbos = run_epoch(model, data, optimizer)
train_elbo, val_elbo = elbos
train_curve.append(train_elbo)
val_curve.append(val_elbo)
print(f"[Epoch {epoch}] train elbo: {train_elbo} val_elbo: {val_elbo}")
# --------------------------------------------------------------------
# Add functionality to plot samples from model during training.
# You can use the make_grid functionality that is already imported.
# --------------------------------------------------------------------
imgs, means = model.sample(n_samples=ARGS.n_samples)
# Only use means for better visualization
# Reshape the data to [n_samples, 1, 28, 28]
means = means.reshape(-1, 1, 28, 28)
save_image(means,
f"grid_Epoch{epoch}.png",
nrow=int(math.sqrt(ARGS.n_samples)),
padding=2,
normalize=True)
# --------------------------------------------------------------------
# Add functionality to plot plot the learned data manifold after
# if required (i.e., if zdim == 2). You can use the make_grid
# functionality that is already imported.
# --------------------------------------------------------------------
if ARGS.zdim == 2:
manifold_means = model.get_manifold(20)
manifold_means = manifold_means.reshape(-1, 1, 28, 28)
save_image(manifold_means,
f"manifold{epoch}.png",
nrow=20,
padding=2,
normalize=True)
save_elbo_plot(train_curve, val_curve, 'elbo.pdf')
def main():
save_folder = "output_vae"
os.makedirs(save_folder, exist_ok=True)
samples_count = 30
data = bmnist()[:2] # ignore test split
model = VAE(hidden_dim=500, z_dim=ARGS.zdim)
optimizer = torch.optim.Adam(model.parameters())
sampled_ims, im_means = model.sample(samples_count)
plot_samples(sampled_ims, im_means, '-1', save_folder)
train_curve, val_curve = [], []
for epoch in range(ARGS.epochs):
train_elbo, val_elbo = run_epoch(model, data, optimizer)
train_curve.append(train_elbo)
val_curve.append(val_elbo)
print(f"[Epoch {epoch+1}] train elbo: {train_elbo} val_elbo: {val_elbo}")
# --------------------------------------------------------------------
# Add functionality to plot samples from model during training.
# You can use the make_grid functionality that is already imported.
# --------------------------------------------------------------------
if epoch in [int(ARGS.epochs/2), ARGS.epochs-1]:
sampled_ims, im_means = model.sample(samples_count)
plot_samples(sampled_ims, im_means, epoch+1, save_folder)
# --------------------------------------------------------------------
# Add functionality to plot plot the learned data manifold after
# if required (i.e., if zdim == 2). You can use the make_grid
# functionality that is already imported.
# --------------------------------------------------------------------
if ARGS.zdim == 2:
steps = 20
inverse_cdf = norm.ppf(np.linspace(1e-3, 1-1e-3, steps))
x, y = np.meshgrid(inverse_cdf, inverse_cdf)
z = torch.tensor([x.flatten(), y.flatten()]).float().t()
manifold = model.decoder(z)
save_image(manifold.view(-1, 1, 28, 28),
os.path.join(save_folder, 'manifold.png'),
nrow=steps, normalize=True)
save_elbo_plot(train_curve, val_curve, os.path.join(save_folder, 'elbo.png'))
def main():
os.makedirs(f'images/vae/{ARGS.zdim}', exist_ok=True)
pin_memory = True if torch.cuda.is_available() else False
data = bmnist(pin_memory=pin_memory)[:2] # ignore test split
model = VAE(z_dim=ARGS.zdim).to(device)
optimizer = torch.optim.Adam(model.parameters())
train_curve, val_curve = [], []
for epoch in range(ARGS.epochs):
elbos = run_epoch(model, data, optimizer)
train_elbo, val_elbo = elbos
train_curve.append(train_elbo)
val_curve.append(val_elbo)
print(f"[Epoch {epoch}] train elbo: {train_elbo} val_elbo: {val_elbo}")
# --------------------------------------------------------------------
# Add functionality to plot samples from model during training.
# You can use the make_grid functioanlity that is already imported.
# --------------------------------------------------------------------
samples, _ = model.sample(n_samples=25)
samples = samples.view(25, 1, 28, 28)
save_image(samples,
os.path.join(
f'images/vae/{ARGS.zdim}',
f'{epoch}_{train_elbo}_{val_elbo}_samples.eps'),
nrow=5,
normalize=True)
if ARGS.zdim == 2:
# Display a 2D manifold of the digits
# Construct grid of latent variable values
grid = np.meshgrid(
norm.ppf(np.linspace(0.00001, 1.0, 10, endpoint=False)),
norm.ppf(np.linspace(0.00001, 1.0, 10, endpoint=False)))
cartesian_grid = torch.FloatTensor(np.array(grid).T.reshape(
(-1, 2))).to(device)
_, means = model.sample(z_samples=cartesian_grid)
save_image(means.view(-1, 1, 28, 28),
os.path.join(f'images/vae/{ARGS.zdim}', f'latent.eps'),
nrow=10,
normalize=True)
save_elbo_plot(train_curve, val_curve, f'elbo_{ARGS.zdim}.pdf')
def main():
data = bmnist()[:2] # ignore test split
model = VAE(z_dim=ARGS.zdim)
optimizer = torch.optim.Adam(model.parameters())
n_samples = 5
train_curve, val_curve = [], []
for epoch in range(ARGS.epochs):
sampled_ims, im_means = model.sample(n_samples)
save_image(make_grid(sampled_ims.view(n_samples, 1, 28, 28)),
ARGS.save_path + 'sample_epoch_{}.png'.format(epoch),
padding=4)
elbos = run_epoch(model, data, optimizer)
train_elbo, val_elbo = elbos
train_curve.append(train_elbo)
val_curve.append(val_elbo)
print("[Epoch {}] train elbo: {} val_elbo: {}".format(
epoch, train_elbo, val_elbo))
# --------------------------------------------------------------------
# Add functionality to plot samples from model during training.
# You can use the make_grid functioanlity that is already imported.
# --------------------------------------------------------------------
if ARGS.zdim == 2:
n = torch.distributions.Normal(torch.zeros(2), torch.ones(2))
latent_variables = [
n.icdf(torch.Tensor([x, y])) for x in torch.arange(0.05, 1, 0.05)
for y in torch.arange(0.05, 1, 0.05)
]
y = model.decoder(torch.stack(latent_variables))
save_image(make_grid(y.view(y.shape[0], 1, 28, 28), nrow=19),
ARGS.save_path + 'manifold.png')
# --------------------------------------------------------------------
# Add functionality to plot plot the learned data manifold after
# if required (i.e., if zdim == 2). You can use the make_grid
# functionality that is already imported.
# --------------------------------------------------------------------
pickle.dump([train_curve, val_curve],
open(ARGS.save_path + 'curves.p', 'wb'))
save_elbo_plot(train_curve, val_curve, ARGS.save_path + 'elbo.pdf')
torch.save_state_dict(ARGS.save_path + 'model.pth')
def main():
data = bmnist()[:2] # ignore test split
model = VAE(z_dim=ARGS.zdim)
optimizer = torch.optim.Adam(model.parameters())
size_width = 28
max_loss = 1000
script_directory = os.path.split(os.path.abspath(__file__))[0]
filepath = 'encoder.model'
encoder_model = os.path.join(script_directory, filepath)
script_directory = os.path.split(os.path.abspath(__file__))[0]
filepath = 'decoder.model'
decoder_model = os.path.join(script_directory, filepath)
min_loss = -1000
train_curve, val_curve = [], []
for epoch in range(ARGS.epochs):
elbos = run_epoch(model, data, optimizer)
train_elbo, val_elbo = elbos
train_curve.append(train_elbo)
val_curve.append(val_elbo)
if min_loss < val_elbo:
print("models saved iter: " + str(epoch))
torch.save(model.encoder, encoder_model)
torch.save(model.decoder, decoder_model)
min_loss = val_elbo
print(f"[Epoch {epoch}] train elbo: {train_elbo} val_elbo: {val_elbo}")
_, mean_sample = model.sample(64)
save_sample(mean_sample, size_width, epoch)
if ARGS.zdim == 2:
print("manifold")
manifold = model.manifold_sample(256)
save_sample(manifold, size_width, epoch, 16)
np.save('curves.npy', {'train': train_curve, 'val': val_curve})
save_elbo_plot(train_curve, val_curve, 'elbo.pdf')
def main():
data = bmnist()[:2] # ignore test split
model = VAE(z_dim=ARGS.zdim)
#print(model)
optimizer = torch.optim.Adam(model.parameters())
# store images in result folder
result_folder = './results'
train_curve, val_curve = [], []
for epoch in range(ARGS.epochs):
n_samples = 25 # allows for 5x5 grid of samples
# plot before training
if epoch == 0:
plot_sample_means(model,
n_samples,
result_folder,
epoch,
name="before_training")
elbos = run_epoch(model, data, optimizer)
train_elbo, val_elbo = elbos
train_curve.append(train_elbo)
val_curve.append(val_elbo)
print(f"[Epoch {epoch}] train elbo: {train_elbo} val_elbo: {val_elbo}")
# --------------------------------------------------------------------
# Add functionality to plot samples from model during training.
# You can use the make_grid functioanlity that is already imported.
# --------------------------------------------------------------------
# samples plotted for each epoch only for latent space = 20
if ARGS.zdim != 2:
plot_sample_means(model, n_samples, result_folder, epoch, name="")
# --------------------------------------------------------------------
# Add functionality to plot plot the learned data manifold after
# if required (i.e., if zdim == 2). You can use the make_grid
# functionality that is already imported.
# --------------------------------------------------------------------
# only for latent dimension = 2 plot manifold
if ARGS.zdim == 2:
plot_manifold(model=model, result_folder=result_folder, epoch=epoch)
save_elbo_plot(train_curve, val_curve, 'elbo.pdf')
def main(ARGS):
data = bmnist()[:2] # ignore test split
#get device
device = get_device(ARGS)
print_args(ARGS)
#initialize model
model = VAE(z_dim=ARGS.zdim, device=device)
print(f"Model loaded: \n{model}")
#initialize optimizer
optimizer = torch.optim.Adam(model.parameters())
print(f"optimizer loaded: \n{optimizer}")
path_model = "./VAE" + ARGS.experiment + "/"
create_dir(path_model)
sample_n_save(model, epoch=0, path_extra=ARGS.experiment)
print(f"first sample saved in {path_model}")
train_curve, dev_curve = [], []
for epoch in range(ARGS.epochs):
elbos = run_epoch(model, data, optimizer)
train_elbo, dev_elbo = elbos
train_curve.append(train_elbo)
dev_curve.append(dev_elbo)
print(f"[Epoch {epoch}] train elbo: {train_elbo} dev_elbo: {dev_elbo}")
#save status after one epoch of training.
sample_n_save(model, epoch=epoch, path_extra=ARGS.experiment)
np.save(path_model + "train_elbo", train_curve)
np.save(path_model + "dev_curve", dev_curve)
torch.save(model.state_dict(),
path_model + model.__class__.__name__ + ".pt")
if ARGS.zdim == 2:
manifold_plot_n_save(model)
save_elbo_plot(train_curve, dev_curve, 'elbo.pdf')
def main():
data = bmnist()[:2] # ignore test split
model = VAE(ARGS.input_dim, hidden_dim=ARGS.hidden_dim, z_dim=ARGS.zdim)
optimizer = torch.optim.Adam(model.parameters())
model.to(device)
train_curve, val_curve = [], []
for epoch in range(ARGS.epochs):
elbos = run_epoch(epoch, model, data, optimizer)
train_elbo, val_elbo = elbos
train_curve.append(train_elbo)
val_curve.append(val_elbo)
print('-------------------------ELBO--------------------------------------')
print(f"[Epoch {epoch}] train elbo: {train_elbo} val_elbo: {val_elbo}")
print('----------------------------------------------------------------')
save_elbo_plot(train_curve, val_curve, 'elbo.pdf')
#generate last reconstruction
model.eval()
with torch.no_grad():
z = torch.randn((100, ARGS.zdim)).to(device)
plt_data = model.decoder(z.float())
comparison = plt_data.view(plt_data.size(0),1,28,28)
save_image(comparison.cpu(), 'results/vae_reconstruction_' + 'final' + '.png', nrow=10)
print("Done")
# --------------------------------------------------------------------
# Add functionality to plot plot the learned data manifold after
# if required (i.e., if zdim == 2). You can use the make_grid
# functionality that is already imported.
# --------------------------------------------------------------------
if ARGS.zdim==2:
print("plot manifold")
model.eval()
with torch.no_grad():
vizManifold(model)
