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)