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