import matplotlib
if not args.use_gui:
matplotlib.use('Agg')
else:
from matplotlib import pyplot as plt
plt.ion()
plt.show()
import os
from dataset import *
from vlae import VLadder
from trainer import NoisyTrainer
if args.gpus is not '':
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpus
dataset = MnistDataset()
model = VLadder(dataset,
name=args.netname,
reg=args.reg,
batch_size=args.batch_size,
restart=args.restart)
trainer = NoisyTrainer(model, dataset, args)
trainer.train()
# TODO remove
#if args.no_train:
# trainer.visualize()
#else:
# trainer.train()
# Log + print error and quit gracefully if there was one
if error:
print(error_msg)
LOG.error(error_msg)
exit(-1)
# Create the dataset object
if args.dataset == 'mnist':
dataset = MnistDataset()
elif args.dataset == 'lsun':
dataset = LSUNDataset(db_path=args.db_path)
elif args.dataset == 'celebA':
dataset = CelebADataset(db_path=args.db_path)
elif args.dataset == 'svhn':
dataset = SVHNDataset(db_path=args.db_path)
else:
LOG.error("Unknown dataset")
exit(-1)
# Construct network and trainer, then let it fly
num_gpus = len(args.gpus.split(','))
model = SequentialVAE(dataset,
name=args.netname,
batch_size=args.batch_size,
logger=LOG,
version=args.version,
base_dir=base_dir,
num_gpus=num_gpus)
trainer = NoisyTrainer(model, dataset, args, LOG, base_dir)
trainer.train()
from trainer import NoisyTrainer
import tensorflow as tf
import numpy as np
tf.set_random_seed(0)
np.random.seed(0)
os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
if args.gpus is not '':
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpus
if args.dataset == 'mnist':
dataset = MnistDataset()
elif args.dataset == 'lsun':
dataset = LSUNDataset(db_path=args.db_path)
elif args.dataset == 'celebA':
dataset = CelebADataset(db_path=args.db_path)
elif args.dataset == 'svhn':
dataset = SVHNDataset(db_path=args.db_path)
else:
print("Unknown dataset")
exit(-1)
model = VLadder(dataset, name=args.netname, reg=args.reg, batch_size=args.batch_size, restart=not args.no_train)
trainer = NoisyTrainer(model, dataset, args)
if args.no_train:
trainer.visualize()
else:
trainer.train()
if not args.use_gui:
matplotlib.use('Agg')
else:
from matplotlib import pyplot as plt
plt.ion()
plt.show()
from dataset import *
from sequential_vae import SequentialVAE
from trainer import NoisyTrainer
if args.gpus is not '':
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpus
if args.dataset == 'mnist':
dataset = MnistDataset()
elif args.dataset == 'lsun':
dataset = LSUNDataset(db_path=args.db_path)
elif args.dataset == 'celebA':
dataset = CelebADataset(db_path=args.db_path)
elif args.dataset == 'svhn':
dataset = SVHNDataset(db_path=args.db_path)
else:
print("Unknown dataset")
exit(-1)
model = SequentialVAE(dataset, name=args.netname, batch_size=args.batch_size)
trainer = NoisyTrainer(model, dataset, args)
trainer.train()