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')
