def main(args): # Check if the output folder is exist if not os.path.exists(args.folder): os.mkdir(args.folder) # Load model model = CVAE().cuda() if torch.cuda.is_available() else CVAE() model.load_state_dict(torch.load(os.path.join(args.folder, 'cvae.pth'))) # Generate sample = torch.randn(args.num, 20) label = torch.from_numpy(np.asarray([args.digits] * args.num)) sample = Variable( sample).cuda() if torch.cuda.is_available() else Variable(sample) sample = model.decode(sample, label).cpu() save_image(sample.view(args.num, 1, 28, 28).data, os.path.join(args.folder, 'generate.png'), nrow=10)
def main(args): # Check if the output folder is exist if not os.path.exists(args.folder): os.mkdir(args.folder) # Load data torch.manual_seed(args.seed) kwargs = {'num_workers': 1, 'pin_memory': True} if args.cuda else {} train_loader = torch.utils.data.DataLoader(datasets.MNIST( './data', train=True, download=True, transform=transforms.ToTensor()), batch_size=args.batch_size, shuffle=True, **kwargs) # Load model model = CVAE().cuda() if torch.cuda.is_available() else CVAE() optimizer = optim.Adam(model.parameters(), lr=1e-3) # Train and generate sample every epoch loss_list = [] for epoch in range(1, args.epochs + 1): model.train() _loss = train(epoch, model, train_loader, optimizer) loss_list.append(_loss) model.eval() sample = torch.randn(100, 20) label = torch.from_numpy(np.asarray(list(range(10)) * 10)) sample = Variable( sample).cuda() if torch.cuda.is_available() else Variable(sample) sample = model.decode(sample, label).cpu() save_image(sample.view(100, 1, 28, 28).data, os.path.join(args.folder, 'sample_' + str(epoch) + '.png'), nrow=10) plt.plot(range(len(loss_list)), loss_list, '-o') plt.savefig(os.path.join(args.folder, 'cvae_loss_curve.png')) torch.save(model.state_dict(), os.path.join(args.folder, 'cvae.pth'))
for i in range(predictions.shape[0]): plt.subplot(4, 4, i + 1) plt.imshow(predictions[i, :, :, 0], cmap='gray') plt.axis('off') file_dir = './image/'+ dire if not os.path.exists(file_dir): os.makedirs(file_dir) plt.savefig(file_dir +'/image_at_epoch_{:04d}.png'.format(degree)) plt.close() def generate_images(model, data): fig = plt.figure(figsize=(4, 4)) for i in range(data.shape[0]): plt.subplot(4, 4, i + 1) plt.imshow(data[i, :, :, 0], cmap='gray') plt.axis('off') plt.show() for i in range(1, 6): model = CVAE(latent_dim=16, beta=i) checkpoint = tf.train.Checkpoint(model=model) checkpoint.restore("checkpoints/2_20method" + str(i) + "/ckpt-10") mean, logvar = model.encode(test_sample) r_m = np.identity(model.latent_dim) z = model.reparameterize(mean, logvar) theta = np.radians(60) c, s = np.cos(theta), np.sin(theta) r_m[0, [0, 1]], r_m[1, [0, 1]] = [c, s], [-s, c] rota_z = matvec(tf.cast(r_m, dtype=tf.float32), z) phi_z = model.decode(rota_z) generate_and_save_images(phi_z, 1, 'test3' + "/beta_test" + str(i))