else:
if not os.path.exists('models'):
os.mkdir('models')
mdl.VAE_trainer(arguments.cuda, train_loader, vis=vis,
EP=EP, config_num=config_num)
torch.save(mdl.state_dict(), path)
#gen_data, seed = mdl.generate_data(100)
#opts = {'title':'VAE generated data config {}'.format(config_num)}
#vis.images(0.5 + 0.5*gen_data.data.cpu(), opts=opts, win='VAE_config_{}'.format(config_num))
if compute_kdes:
num_samples = 1
scores_iml, gen_data = get_scores(test_loader, mdl, arguments.cuda,
num_samples=num_samples,
task='celeba', base_dist='fixed_iso_gauss')
vis.images(gen_data.data*0.5 + 0.5, win='VAE genim')
pdb.set_trace()
#vis.image(im_gen*0.5 + 0.5, win='VAE genim')
#vis.image(im_test*0.5 + 0.5, win='VAE testim')
elif model == 'VaDE':
EP = 5 # additional iterations
num_samples = 3
Kss = [[100, 100]]
L = 64*64*3
use_gmms = [1, 1, 0]
compute_iml_scores = False
else:
if not os.path.exists('models'):
os.mkdir('models')
# train the VAE
mdl.VAE_trainer(arguments.cuda,
train_loader,
vis=vis,
EP=EP,
config_num=config_num)
torch.save(mdl.state_dict(), path)
scores_vae, gen_data = get_scores(test_loader,
mdl,
arguments.cuda,
num_samples=num_samples,
task='mnist',
base_dist='fixed_iso_gauss')
elif model == 'VaDE':
EP = 100
num_samples = 1
for config_num, Ks in enumerate(Kss):
print(Ks)
mdl = VAE(L1, L2, Ks, M=M)
if arguments.cuda:
mdl.cuda()
arguments = args()
arguments.batch_size = 3000
arguments.data = 'mnist'
arguments.input_type = 'autoenc'
arguments.cuda = True
dataset_loader, test_loader = ut.get_loaders(arguments.batch_size, c=0,
arguments=arguments)
# compute the scores here for plain GAN
num_samples = 1
if 1:
print('Computing WGAN scores..')
scores_gan, _ = get_scores(test_loader, dec, arguments.cuda,
num_samples=num_samples,
task='mnist', base_dist='fixed_iso_gauss')
enc = encoder(784, [40, 600])
if cuda:
enc = enc.cuda()
opt = optim.Adam(enc.parameters(), lr=1e-4, betas=(0.5, 0.9))
if 0 and os.path.exists('mnist_encoder.pt'):
enc.load_state_dict(torch.load('mnist_encoder.pt'))
else:
for ep in range(100):
for i, (tar, _) in enumerate(dataset_loader):
if cuda:
tar = tar.cuda()