def main(_):
pp.pprint(vars(FLAGS))
sess_config = tf.ConfigProto(
device_count={"CPU": 3},
inter_op_parallelism_threads=0,
intra_op_parallelism_threads=0,
allow_soft_placement=True)
sess_config.gpu_options.allow_growth = True
if FLAGS.model == 'mmd':
from core.model import MMD_GAN as Model
elif FLAGS.model == 'gan':
from core.gan import GAN as Model
elif FLAGS.model == 'wgan_gp':
from core.wgan_gp import WGAN_GP as Model
elif FLAGS.model == 'cramer':
from core.cramer import Cramer_GAN as Model
elif FLAGS.model == 'smmd':
from core.smmd import SMMD as Model
elif FLAGS.model == 'swgan':
from core.smmd import SWGAN as Model
else:
raise ValueError("unknown model {}".format(FLAGS.model))
#if FLAGS.multi_gpu:
# from core.model_multi_gpu import MMD_GAN as Model
with tf.Session(config=sess_config) as sess:
#sess = tf_debug.tf_debug.TensorBoardDebugWrapperSession(sess,'localhost:6064')
#sess.add_tensor_filter("has_inf_or_nan", tf_debug.has_inf_or_nan)
if FLAGS.dataset == 'mnist':
gan = Model(sess, config=FLAGS, batch_size=FLAGS.batch_size, output_size=28, c_dim=1,
data_dir=FLAGS.data_dir)
elif FLAGS.dataset == 'cifar10':
gan = Model(sess, config=FLAGS, batch_size=FLAGS.batch_size, output_size=32, c_dim=3,
data_dir=FLAGS.data_dir)
elif FLAGS.dataset in ['celebA', 'lsun', 'imagenet']:
gan = Model(sess, config=FLAGS, batch_size=FLAGS.batch_size, output_size=FLAGS.output_size, c_dim=3,
data_dir=FLAGS.data_dir)
else:
gan = Model(
sess, batch_size=FLAGS.batch_size,
output_size=FLAGS.output_size, c_dim=FLAGS.c_dim,
data_dir=FLAGS.data_dir)
if FLAGS.is_train:
gan.train()
gan.pre_process_only()
elif FLAGS.print_pca:
gan.print_pca()
elif FLAGS.visualize:
gan.load_checkpoint()
visualize(sess, gan, FLAGS, 2)
else:
gan.get_samples(FLAGS.no_of_samples, layers=[-1])
if FLAGS.log:
sys.stdout = gan.old_stdout
gan.log_file.close()
gan.sess.close()
def main(_):
pp.pprint(FLAGS.__flags)
if FLAGS.threads < np.inf:
sess_config = tf.ConfigProto(
intra_op_parallelism_threads=FLAGS.threads)
sess_config.gpu_options.per_process_gpu_memory_fraction = FLAGS.gpu_mem
else:
sess_config = tf.ConfigProto()
if 'mmd' in FLAGS.model:
from core.model import MMD_GAN as Model
elif FLAGS.model == 'wgan_gp':
from core.wgan_gp import WGAN_GP as Model
elif 'cramer' in FLAGS.model:
from core.cramer import Cramer_GAN as Model
with tf.Session(config=sess_config) as sess:
if FLAGS.dataset == 'mnist':
gan = Model(sess,
config=FLAGS,
batch_size=FLAGS.batch_size,
output_size=28,
c_dim=1,
data_dir=FLAGS.data_dir)
elif FLAGS.dataset == 'cifar10':
gan = Model(sess,
config=FLAGS,
batch_size=FLAGS.batch_size,
output_size=32,
c_dim=3,
data_dir=FLAGS.data_dir)
elif FLAGS.dataset in ['celebA', 'lsun']:
gan = Model(sess,
config=FLAGS,
batch_size=FLAGS.batch_size,
output_size=FLAGS.output_size,
c_dim=3,
data_dir=FLAGS.data_dir)
else:
gan = Model(sess,
batch_size=FLAGS.batch_size,
output_size=FLAGS.output_size,
c_dim=FLAGS.c_dim,
data_dir=FLAGS.data_dir)
if FLAGS.is_train:
gan.train()
elif FLAGS.print_pca:
gan.print_pca()
elif FLAGS.visualize:
gan.load_checkpoint()
visualize(sess, gan, FLAGS, 2)
else:
gan.get_samples(FLAGS.no_of_samples, layers=[-1])
if FLAGS.log:
sys.stdout = gan.old_stdout
gan.log_file.close()
gan.sess.close()
def main(_):
global FLAGS
pp.pprint(vars(FLAGS))
sess_config = tf.ConfigProto(device_count={"CPU": 3},
inter_op_parallelism_threads=0,
intra_op_parallelism_threads=0,
allow_soft_placement=True)
sess_config.gpu_options.allow_growth = True
if FLAGS.dataset == 'mnist':
FLAGS.output_size = 28
FLAGS.c_dim = 1
elif FLAGS.dataset == 'cifar10':
FLAGS.output_size = 32
FLAGS.c_dim = 3
elif FLAGS.dataset in ['celebA', 'lsun', 'imagenet']:
FLAGS.c_dim = 3
from core import model_class
Model = model_class(FLAGS.model)
with tf.Session(config=sess_config) as sess:
#sess = tf_debug.tf_debug.TensorBoardDebugWrapperSession(sess,'localhost:6064')
#sess.add_tensor_filter("has_inf_or_nan", tf_debug.has_inf_or_nan)
gan = Model(sess, config=FLAGS)
if FLAGS.is_train:
gan.train()
elif FLAGS.print_pca:
gan.print_pca()
elif FLAGS.visualize:
gan.load_checkpoint()
visualize(sess, gan, FLAGS, 2)
else:
gan.get_samples(FLAGS.no_of_samples, layers=[-1])
if FLAGS.log:
sys.stdout = gan.old_stdout
gan.log_file.close()
gan.sess.close()
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--batch_size',
type=int,
default=32,
help='Mini-batch size.')
parser.add_argument('--input_dim',
type=int,
default=257,
help='The dimension of inputs.')
parser.add_argument('--output_dim',
type=int,
default=40,
help='The dimension of outputs.')
parser.add_argument('--num_threads',
type=int,
default=1,
help='The num of threads to read tfrecords files.')
parser.add_argument('--num_epochs',
type=int,
default=1,
help='The num of epochs to read tfrecords files.')
parser.add_argument('--data_dir',
type=str,
default='data/tfrecords/',
help='Directory of train, val and test data.')
FLAGS, unparsed = parser.parse_known_args()
pp.pprint(FLAGS.__dict__)
sys.stdout.flush()
main()
