model = Transporter(encodings, DISTR, FOLDER, BATCH_SIZE_GEN) elif MODEL == 'generator': model = Generator(encodings, DISTR, FOLDER, BATCH_SIZE_GEN) # I Could try L2 Loss instead of L1? else: raise NotImplementedError # Train the Latent Space Model if GEN_LOAD: model.load_weights(MODEL) else: model.train(STEPS, lr=0.001) # I should try adjusting the learning rate? #model.train(STEPS//2, lr=0.0003) #model.train(STEPS//2, lr=0.0001) # Display Results fake_distr = model.generate(batches=1) fake_img = ae.decode(fake_distr) fake_img = np.reshape(fake_img, ((BATCH_SIZE_GEN,) + shape)) save_image(torch.Tensor(fake_img[0:64]), os.path.join(FOLDER, "final.png")) # Save Images in a file for later evaluation eval_distr = model.generate(batches=10000//BATCH_SIZE_GEN) eval_img = ae.decode(eval_distr) np.save("{}/distribution.npy".format(FOLDER), eval_distr) np.save("{}/images.npy".format(FOLDER), eval_img) #channels_last = np.rollaxis(fake_img[0:16], 1, 4) #display_img(channels_last, columns=4)
latent, test = make_moons() elif DATASET == 'two_cluster': latent, test = two_cluster() elif DATASET == 'eight_cluster': latent, test = eight_cluster() elif DATASET == 'circles': latent, test = make_circles() else: raise NotImplementedError # Prepare the Latent Space Model if MODEL == 'transporter': model = Transporter(latent, DISTR, FOLDER, BATCH_SIZE_GEN) elif MODEL == 'generator': model = Generator(latent, DISTR, FOLDER, BATCH_SIZE_GEN) else: raise NotImplementedError # Train the Latent Space Model if GEN_LOAD: model.load_weights(MODEL) else: model.train(STEPS, lr=0.0001, images=True) # Evaluate model.evaluate() # Display Results fake_distr = model.generate(batches=1) save_points(fake_distr, latent[0:BATCH_SIZE_GEN], FOLDER, name='final')