def main(args):
# clear param store
pyro.clear_param_store()
# setup MNIST data loaders
# train_loader, test_loader
train_loader, test_loader = setup_data_loaders(MNIST,
use_cuda=args.cuda,
batch_size=256)
# setup the VAE
vae = VAE(use_cuda=args.cuda)
# setup the optimizer
adam_args = {"lr": args.learning_rate}
optimizer = Adam(adam_args)
# setup the inference algorithm
elbo = JitTrace_ELBO() if args.jit else Trace_ELBO()
svi = SVI(vae.model, vae.guide, optimizer, loss=elbo)
# setup visdom for visualization
if args.visdom_flag:
vis = visdom.Visdom()
train_elbo = []
test_elbo = []
# training loop
for epoch in range(args.num_epochs):
# initialize loss accumulator
epoch_loss = 0.
# do a training epoch over each mini-batch x returned
# by the data loader
for x, _ in train_loader:
# if on GPU put mini-batch into CUDA memory
if args.cuda:
x = x.cuda()
# do ELBO gradient and accumulate loss
epoch_loss += svi.step(x)
# report training diagnostics
normalizer_train = len(train_loader.dataset)
total_epoch_loss_train = epoch_loss / normalizer_train
train_elbo.append(total_epoch_loss_train)
print("[epoch %03d] average training loss: %.4f" %
(epoch, total_epoch_loss_train))
if epoch % args.test_frequency == 0:
# initialize loss accumulator
test_loss = 0.
# compute the loss over the entire test set
for i, (x, _) in enumerate(test_loader):
# if on GPU put mini-batch into CUDA memory
if args.cuda:
x = x.cuda()
# compute ELBO estimate and accumulate loss
test_loss += svi.evaluate_loss(x)
# pick three random test images from the first mini-batch and
# visualize how well we're reconstructing them
if i == 0:
if args.visdom_flag:
plot_vae_samples(vae, vis)
reco_indices = np.random.randint(0, x.shape[0], 3)
for index in reco_indices:
test_img = x[index, :]
reco_img = vae.reconstruct_img(test_img)
vis.image(test_img.reshape(
28, 28).detach().cpu().numpy(),
opts={'caption': 'test image'})
vis.image(reco_img.reshape(
28, 28).detach().cpu().numpy(),
opts={'caption': 'reconstructed image'})
# report test diagnostics
normalizer_test = len(test_loader.dataset)
total_epoch_loss_test = test_loss / normalizer_test
test_elbo.append(total_epoch_loss_test)
print("[epoch %03d] average test loss: %.4f" %
(epoch, total_epoch_loss_test))
if epoch == args.tsne_iter:
mnist_test_tsne(vae=vae, test_loader=test_loader)
plot_llk(np.array(train_elbo), np.array(test_elbo))
return vae
def main(args):
"""
run inference for SS-VAE
:param args: arguments for SS-VAE
:return: None
"""
if args.seed is not None:
pyro.set_rng_seed(args.seed)
# batch_size: number of images (and labels) to be considered in a batch
ss_vae = SSVAE(z_dim=args.z_dim,
hidden_layers=args.hidden_layers,
use_cuda=args.cuda,
config_enum=args.enum_discrete,
aux_loss_multiplier=args.aux_loss_multiplier)
# setup the optimizer
adam_params = {"lr": args.learning_rate, "betas": (args.beta_1, 0.999)}
optimizer = Adam(adam_params)
# set up the loss(es) for inference. wrapping the guide in config_enumerate builds the loss as a sum
# by enumerating each class label for the sampled discrete categorical distribution in the model
guide_enum = config_enumerate(guide, args.enum_discrete, expand=True)
elbo = (JitTraceEnum_ELBO if args.jit else TraceEnum_ELBO)(max_plate_nesting=1)
loss_basic = SVI(model, guide_enum, optimizer, loss=elbo)
# build a list of all losses considered
losses = [loss_basic]
# aux_loss: whether to use the auxiliary loss from NIPS 14 paper (Kingma et al)
if args.aux_loss:
elbo = JitTrace_ELBO() if args.jit else Trace_ELBO()
loss_aux = SVI(ss_vae.model_classify, ss_vae.guide_classify, optimizer, loss=elbo)
losses.append(loss_aux)
try:
# setup the logger if a filename is provided
logger = open(args.logfile, "w") if args.logfile else None
data_loaders = setup_data_loaders(MNISTCached, args.cuda, args.batch_size, sup_num=args.sup_num)
# how often would a supervised batch be encountered during inference
# e.g. if sup_num is 3000, we would have every 16th = int(50000/3000) batch supervised
# until we have traversed through the all supervised batches
periodic_interval_batches = int(MNISTCached.train_data_size / (1.0 * args.sup_num))
# number of unsupervised examples
unsup_num = MNISTCached.train_data_size - args.sup_num
# initializing local variables to maintain the best validation accuracy
# seen across epochs over the supervised training set
# and the corresponding testing set and the state of the networks
best_valid_acc, corresponding_test_acc = 0.0, 0.0
# WL: added. =====
print_and_log(logger, args)
print_and_log(logger, "\nepoch\t"+"elbo(sup)\t"+"elbo(unsup)\t"+"time(sec)")
times = [time.time()]
# ================
# run inference for a certain number of epochs
for i in range(0, args.num_epochs):
# get the losses for an epoch
epoch_losses_sup, epoch_losses_unsup = \
run_inference_for_epoch(data_loaders, losses, periodic_interval_batches)
# compute average epoch losses i.e. losses per example
avg_epoch_losses_sup = map(lambda v: v / args.sup_num, epoch_losses_sup)
avg_epoch_losses_unsup = map(lambda v: v / unsup_num, epoch_losses_unsup)
# store the loss and validation/testing accuracies in the logfile
# WL: edited. =====
# str_loss_sup = " ".join(map(str, avg_epoch_losses_sup))
# str_loss_unsup = " ".join(map(str, avg_epoch_losses_unsup))
# str_print = "{} epoch: avg losses {}".format(i, "{} {}".format(str_loss_sup, str_loss_unsup))
times.append(time.time())
str_elbo_sup = " ".join(map(lambda v: f"{-v:.4f}", avg_epoch_losses_sup))
str_elbo_unsup = " ".join(map(lambda v: f"{-v:.4f}", avg_epoch_losses_unsup))
str_print = f"{i:06d}\t"\
f"{str_elbo_sup}\t"\
f"{str_elbo_unsup}\t"\
f"{times[-1]-times[-2]:.3f}"
# =================
validation_accuracy = get_accuracy(data_loaders["valid"], ss_vae.classifier, args.batch_size)
# WL: commented. =====
# str_print += " validation accuracy {}".format(validation_accuracy)
# ====================
# this test accuracy is only for logging, this is not used
# to make any decisions during training
test_accuracy = get_accuracy(data_loaders["test"], ss_vae.classifier, args.batch_size)
# WL: commented. =====
# str_print += " test accuracy {}".format(test_accuracy)
# ====================
# update the best validation accuracy and the corresponding
# testing accuracy and the state of the parent module (including the networks)
if best_valid_acc < validation_accuracy:
best_valid_acc = validation_accuracy
corresponding_test_acc = test_accuracy
print_and_log(logger, str_print)
final_test_accuracy = get_accuracy(data_loaders["test"], ss_vae.classifier, args.batch_size)
# WL: commented. =====
# print_and_log(logger, "best validation accuracy {} corresponding testing accuracy {} "
# "last testing accuracy {}".format(best_valid_acc, corresponding_test_acc, final_test_accuracy))
# ====================
finally:
# close the logger file object if we opened it earlier
if args.logfile:
logger.close()
def main():
# parse command line arguments
parser = argparse.ArgumentParser(description="parse args")
parser.add_argument('-n',
'--num-epochs',
default=101,
type=int,
help='number of training epochs')
parser.add_argument('-tf',
'--test-frequency',
default=5,
type=int,
help='how often we evaluate the test set')
parser.add_argument('-lr',
'--learning-rate',
default=1.0e-3,
type=float,
help='learning rate')
parser.add_argument('-b1',
'--beta1',
default=0.95,
type=float,
help='beta1 adam hyperparameter')
parser.add_argument('--cuda',
action='store_true',
default=False,
help='whether to use cuda')
parser.add_argument('-visdom',
'--visdom_flag',
default=False,
help='Whether plotting in visdom is desired')
parser.add_argument('-i-tsne',
'--tsne_iter',
default=100,
type=int,
help='epoch when tsne visualization runs')
args = parser.parse_args()
# setup MNIST data loaders
# train_loader, test_loader
train_loader, test_loader = setup_data_loaders(MNIST,
use_cuda=args.cuda,
batch_size=256)
# setup the VAE
vae = VAE(use_cuda=args.cuda)
# setup the optimizer
adam_args = {"lr": args.learning_rate}
optimizer = Adam(adam_args)
# setup the inference algorithm
svi = SVI(vae.model, vae.guide, optimizer, loss="ELBO")
# setup visdom for visualization
if args.visdom_flag:
vis = visdom.Visdom()
train_elbo = []
test_elbo = []
# training loop
for epoch in range(args.num_epochs):
# initialize loss accumulator
epoch_loss = 0.
# do a training epoch over each mini-batch x returned
# by the data loader
for _, (x, _) in enumerate(train_loader):
# if on GPU put mini-batch into CUDA memory
if args.cuda:
x = x.cuda()
# wrap the mini-batch in a PyTorch Variable
x = Variable(x)
# do ELBO gradient and accumulate loss
epoch_loss += svi.step(x)
# report training diagnostics
normalizer_train = len(train_loader.dataset)
total_epoch_loss_train = epoch_loss / normalizer_train
train_elbo.append(total_epoch_loss_train)
print("[epoch %03d] average training loss: %.4f" %
(epoch, total_epoch_loss_train))
if epoch % args.test_frequency == 0:
# initialize loss accumulator
test_loss = 0.
# compute the loss over the entire test set
for i, (x, _) in enumerate(test_loader):
# if on GPU put mini-batch into CUDA memory
if args.cuda:
x = x.cuda()
# wrap the mini-batch in a PyTorch Variable
x = Variable(x)
# compute ELBO estimate and accumulate loss
test_loss += svi.evaluate_loss(x)
# pick three random test images from the first mini-batch and
# visualize how well we're reconstructing them
if i == 0:
if args.visdom_flag:
plot_vae_samples(vae, vis)
reco_indices = np.random.randint(0, x.size(0), 3)
for index in reco_indices:
test_img = x[index, :]
reco_img = vae.reconstruct_img(test_img)
vis.image(test_img.contiguous().view(
28, 28).data.cpu().numpy(),
opts={'caption': 'test image'})
vis.image(reco_img.contiguous().view(
28, 28).data.cpu().numpy(),
opts={'caption': 'reconstructed image'})
# report test diagnostics
normalizer_test = len(test_loader.dataset)
total_epoch_loss_test = test_loss / normalizer_test
test_elbo.append(total_epoch_loss_test)
print("[epoch %03d] average test loss: %.4f" %
(epoch, total_epoch_loss_test))
if epoch == args.tsne_iter:
mnist_test_tsne(vae=vae, test_loader=test_loader)
plot_llk(np.array(train_elbo), np.array(test_elbo))
return vae
def run_inference_ss_vae(args):
"""
run inference for SS-VAE
:param args: arguments for SS-VAE
:return: None
"""
if args.use_cuda:
torch.set_default_tensor_type('torch.cuda.FloatTensor')
if args.seed is not None:
set_seed(args.seed, args.use_cuda)
viz = None
if args.visualize:
from visdom import Visdom
viz = Visdom()
mkdir_p("./vae_results")
# batch_size: number of images (and labels) to be considered in a batch
ss_vae = SSVAE(z_dim=args.z_dim,
hidden_layers=args.hidden_layers,
epsilon_scale=args.epsilon_scale,
use_cuda=args.use_cuda,
aux_loss_multiplier=args.aux_loss_multiplier)
# setup the optimizer
adam_params = {"lr": args.learning_rate, "betas": (args.beta_1, 0.999)}
optimizer = Adam(adam_params)
# set up the loss(es) for inference setting the enum_discrete parameter builds the loss as a sum
# by enumerating each class label for the sampled discrete categorical distribution in the model
loss_basic = SVI(ss_vae.model, ss_vae.guide, optimizer, loss="ELBO", enum_discrete=args.enum_discrete)
# build a list of all losses considered
losses = [loss_basic]
# aux_loss: whether to use the auxiliary loss from NIPS 14 paper (Kingma et al)
if args.aux_loss:
loss_aux = SVI(ss_vae.model_classify, ss_vae.guide_classify, optimizer, loss="ELBO")
losses.append(loss_aux)
try:
# setup the logger if a filename is provided
logger = None if args.logfile is None else open(args.logfile, "w")
data_loaders = setup_data_loaders(MNISTCached, args.use_cuda, args.batch_size, sup_num=args.sup_num)
# how often would a supervised batch be encountered during inference
# e.g. if sup_num is 3000, we would have every 16th = int(50000/3000) batch supervised
# until we have traversed through the all supervised batches
periodic_interval_batches = int(MNISTCached.train_data_size / (1.0 * args.sup_num))
# number of unsupervised examples
unsup_num = MNISTCached.train_data_size - args.sup_num
# initializing local variables to maintain the best validation accuracy
# seen across epochs over the supervised training set
# and the corresponding testing set and the state of the networks
best_valid_acc, corresponding_test_acc = 0.0, 0.0
# run inference for a certain number of epochs
for i in range(0, args.num_epochs):
# get the losses for an epoch
epoch_losses_sup, epoch_losses_unsup = \
run_inference_for_epoch(data_loaders, losses, periodic_interval_batches)
# compute average epoch losses i.e. losses per example
avg_epoch_losses_sup = map(lambda v: v / args.sup_num, epoch_losses_sup)
avg_epoch_losses_unsup = map(lambda v: v / unsup_num, epoch_losses_unsup)
# store the loss and validation/testing accuracies in the logfile
str_loss_sup = " ".join(map(str, avg_epoch_losses_sup))
str_loss_unsup = " ".join(map(str, avg_epoch_losses_unsup))
str_print = "{} epoch: avg losses {}".format(i, "{} {}".format(str_loss_sup, str_loss_unsup))
validation_accuracy = get_accuracy(data_loaders["valid"], ss_vae.classifier, args.batch_size)
str_print += " validation accuracy {}".format(validation_accuracy)
# this test accuracy is only for logging, this is not used
# to make any decisions during training
test_accuracy = get_accuracy(data_loaders["test"], ss_vae.classifier, args.batch_size)
str_print += " test accuracy {}".format(test_accuracy)
# update the best validation accuracy and the corresponding
# testing accuracy and the state of the parent module (including the networks)
if best_valid_acc < validation_accuracy:
best_valid_acc = validation_accuracy
corresponding_test_acc = test_accuracy
print_and_log(logger, str_print)
final_test_accuracy = get_accuracy(data_loaders["test"], ss_vae.classifier, args.batch_size)
print_and_log(logger, "best validation accuracy {} corresponding testing accuracy {} "
"last testing accuracy {}".format(best_valid_acc, corresponding_test_acc, final_test_accuracy))
# visualize the conditional samples
visualize(ss_vae, viz, data_loaders["test"])
finally:
# close the logger file object if we opened it earlier
if args.logfile is not None:
logger.close()
def main():
# parse command line arguments
parser = argparse.ArgumentParser(description="parse args")
parser.add_argument('-n', '--num-epochs', default=101, type=int, help='number of training epochs')
parser.add_argument('-tf', '--test-frequency', default=5, type=int, help='how often we evaluate the test set')
parser.add_argument('-lr', '--learning-rate', default=1.0e-3, type=float, help='learning rate')
parser.add_argument('-b1', '--beta1', default=0.95, type=float, help='beta1 adam hyperparameter')
parser.add_argument('--cuda', action='store_true', default=False, help='whether to use cuda')
parser.add_argument('-visdom', '--visdom_flag', default=False, help='Whether plotting in visdom is desired')
parser.add_argument('-i-tsne', '--tsne_iter', default=100, type=int, help='epoch when tsne visualization runs')
args = parser.parse_args()
# setup MNIST data loaders
# train_loader, test_loader
train_loader, test_loader = setup_data_loaders(MNIST, use_cuda=args.cuda, batch_size=256)
# setup the VAE
vae = VAE(use_cuda=args.cuda)
# setup the optimizer
adam_args = {"lr": args.learning_rate}
optimizer = Adam(adam_args)
# setup the inference algorithm
svi = SVI(vae.model, vae.guide, optimizer, loss="ELBO")
# setup visdom for visualization
if args.visdom_flag:
vis = visdom.Visdom()
train_elbo = []
test_elbo = []
# training loop
for epoch in range(args.num_epochs):
# initialize loss accumulator
epoch_loss = 0.
# do a training epoch over each mini-batch x returned
# by the data loader
for _, (x, _) in enumerate(train_loader):
# if on GPU put mini-batch into CUDA memory
if args.cuda:
x = x.cuda()
# wrap the mini-batch in a PyTorch Variable
x = Variable(x)
# do ELBO gradient and accumulate loss
epoch_loss += svi.step(x)
# report training diagnostics
normalizer_train = len(train_loader.dataset)
total_epoch_loss_train = epoch_loss / normalizer_train
train_elbo.append(total_epoch_loss_train)
print("[epoch %03d] average training loss: %.4f" % (epoch, total_epoch_loss_train))
if epoch % args.test_frequency == 0:
# initialize loss accumulator
test_loss = 0.
# compute the loss over the entire test set
for i, (x, _) in enumerate(test_loader):
# if on GPU put mini-batch into CUDA memory
if args.cuda:
x = x.cuda()
# wrap the mini-batch in a PyTorch Variable
x = Variable(x)
# compute ELBO estimate and accumulate loss
test_loss += svi.evaluate_loss(x)
# pick three random test images from the first mini-batch and
# visualize how well we're reconstructing them
if i == 0:
if args.visdom_flag:
plot_vae_samples(vae, vis)
reco_indices = np.random.randint(0, x.size(0), 3)
for index in reco_indices:
test_img = x[index, :]
reco_img = vae.reconstruct_img(test_img)
vis.image(test_img.contiguous().view(28, 28).data.cpu().numpy(),
opts={'caption': 'test image'})
vis.image(reco_img.contiguous().view(28, 28).data.cpu().numpy(),
opts={'caption': 'reconstructed image'})
# report test diagnostics
normalizer_test = len(test_loader.dataset)
total_epoch_loss_test = test_loss / normalizer_test
test_elbo.append(total_epoch_loss_test)
print("[epoch %03d] average test loss: %.4f" % (epoch, total_epoch_loss_test))
if epoch == args.tsne_iter:
mnist_test_tsne(vae=vae, test_loader=test_loader)
plot_llk(np.array(train_elbo), np.array(test_elbo))
return vae
def run_inference_ss_vae(args):
"""
run inference for SS-VAE
:param args: arguments for SS-VAE
:return: None
"""
if args.use_cuda:
torch.set_default_tensor_type('torch.cuda.FloatTensor')
if args.seed is not None:
set_seed(args.seed, args.use_cuda)
viz = None
if args.visualize:
from visdom import Visdom
viz = Visdom()
mkdir_p("./vae_results")
# batch_size: number of images (and labels) to be considered in a batch
ss_vae = SSVAE(z_dim=args.z_dim,
hidden_layers=args.hidden_layers,
epsilon_scale=args.epsilon_scale,
use_cuda=args.use_cuda,
aux_loss_multiplier=args.aux_loss_multiplier)
# setup the optimizer
adam_params = {"lr": args.learning_rate, "betas": (args.beta_1, 0.999)}
optimizer = Adam(adam_params)
# set up the loss(es) for inference setting the enum_discrete parameter builds the loss as a sum
# by enumerating each class label for the sampled discrete categorical distribution in the model
loss_basic = SVI(ss_vae.model, ss_vae.guide, optimizer, loss="ELBO", enum_discrete=args.enum_discrete)
# build a list of all losses considered
losses = [loss_basic]
# aux_loss: whether to use the auxiliary loss from NIPS 14 paper (Kingma et al)
if args.aux_loss:
loss_aux = SVI(ss_vae.model_classify, ss_vae.guide_classify, optimizer, loss="ELBO")
losses.append(loss_aux)
try:
# setup the logger if a filename is provided
logger = None if args.logfile is None else open(args.logfile, "w")
data_loaders = setup_data_loaders(MNISTCached, args.use_cuda, args.batch_size, sup_num=args.sup_num)
# how often would a supervised batch be encountered during inference
# e.g. if sup_num is 3000, we would have every 16th = int(50000/3000) batch supervised
# until we have traversed through the all supervised batches
periodic_interval_batches = int(MNISTCached.train_data_size / (1.0 * args.sup_num))
# number of unsupervised examples
unsup_num = MNISTCached.train_data_size - args.sup_num
# initializing local variables to maintain the best validation accuracy
# seen across epochs over the supervised training set
# and the corresponding testing set and the state of the networks
best_valid_acc, corresponding_test_acc = 0.0, 0.0
# run inference for a certain number of epochs
for i in range(0, args.num_epochs):
# get the losses for an epoch
epoch_losses_sup, epoch_losses_unsup = \
run_inference_for_epoch(data_loaders, losses, periodic_interval_batches)
# compute average epoch losses i.e. losses per example
avg_epoch_losses_sup = map(lambda v: v / args.sup_num, epoch_losses_sup)
avg_epoch_losses_unsup = map(lambda v: v / unsup_num, epoch_losses_unsup)
# store the loss and validation/testing accuracies in the logfile
str_loss_sup = " ".join(map(str, avg_epoch_losses_sup))
str_loss_unsup = " ".join(map(str, avg_epoch_losses_unsup))
str_print = "{} epoch: avg losses {}".format(i, "{} {}".format(str_loss_sup, str_loss_unsup))
validation_accuracy = get_accuracy(data_loaders["valid"], ss_vae.classifier, args.batch_size)
str_print += " validation accuracy {}".format(validation_accuracy)
# this test accuracy is only for logging, this is not used
# to make any decisions during training
test_accuracy = get_accuracy(data_loaders["test"], ss_vae.classifier, args.batch_size)
str_print += " test accuracy {}".format(test_accuracy)
# update the best validation accuracy and the corresponding
# testing accuracy and the state of the parent module (including the networks)
if best_valid_acc < validation_accuracy:
best_valid_acc = validation_accuracy
corresponding_test_acc = test_accuracy
print_and_log(logger, str_print)
final_test_accuracy = get_accuracy(data_loaders["test"], ss_vae.classifier, args.batch_size)
print_and_log(logger, "best validation accuracy {} corresponding testing accuracy {} "
"last testing accuracy {}".format(best_valid_acc, corresponding_test_acc, final_test_accuracy))
# visualize the conditional samples
visualize(ss_vae, viz, data_loaders["test"])
finally:
# close the logger file object if we opened it earlier
if args.logfile is not None:
logger.close()
def main(args):
# setup MNIST data loaders
# train_loader, test_loader
train_loader, test_loader = setup_data_loaders(MNIST, use_cuda=args.cuda, batch_size=256)
# setup the VAE
vae = VAE(use_cuda=args.cuda)
# setup the optimizer
adam_args = {"lr": args.learning_rate}
optimizer = Adam(adam_args)
# setup the inference algorithm
svi = SVI(vae.model, vae.guide, optimizer, loss=Trace_ELBO())
# setup visdom for visualization
if args.visdom_flag:
vis = visdom.Visdom()
train_elbo = []
test_elbo = []
# training loop
for epoch in range(args.num_epochs):
# initialize loss accumulator
epoch_loss = 0.
# do a training epoch over each mini-batch x returned
# by the data loader
for _, (x, _) in enumerate(train_loader):
# if on GPU put mini-batch into CUDA memory
if args.cuda:
x = x.cuda()
# do ELBO gradient and accumulate loss
epoch_loss += svi.step(x)
# report training diagnostics
normalizer_train = len(train_loader.dataset)
total_epoch_loss_train = epoch_loss / normalizer_train
train_elbo.append(total_epoch_loss_train)
print("[epoch %03d] average training loss: %.4f" % (epoch, total_epoch_loss_train))
if epoch % args.test_frequency == 0:
# initialize loss accumulator
test_loss = 0.
# compute the loss over the entire test set
for i, (x, _) in enumerate(test_loader):
# if on GPU put mini-batch into CUDA memory
if args.cuda:
x = x.cuda()
# compute ELBO estimate and accumulate loss
test_loss += svi.evaluate_loss(x)
# pick three random test images from the first mini-batch and
# visualize how well we're reconstructing them
if i == 0:
if args.visdom_flag:
plot_vae_samples(vae, vis)
reco_indices = np.random.randint(0, x.size(0), 3)
for index in reco_indices:
test_img = x[index, :]
reco_img = vae.reconstruct_img(test_img)
vis.image(test_img.reshape(28, 28).detach().cpu().numpy(),
opts={'caption': 'test image'})
vis.image(reco_img.reshape(28, 28).detach().cpu().numpy(),
opts={'caption': 'reconstructed image'})
# report test diagnostics
normalizer_test = len(test_loader.dataset)
total_epoch_loss_test = test_loss / normalizer_test
test_elbo.append(total_epoch_loss_test)
print("[epoch %03d] average test loss: %.4f" % (epoch, total_epoch_loss_test))
if epoch == args.tsne_iter:
mnist_test_tsne(vae=vae, test_loader=test_loader)
plot_llk(np.array(train_elbo), np.array(test_elbo))
return vae
