def main(args): model = None if args.model == 'GAN': model = GAN(args) elif args.model == 'DCGAN': model = DCGAN_MODEL(args) elif args.model == 'WGAN-CP': model = WGAN_CP(args) elif args.model == 'WGAN-GP': model = WGAN_GP(args) elif args.model == 'VLGAN': model = VLGAN_MODEL(args) else: print("Model type non-existing. Try again.") exit(-1) # Load datasets to train and test loaders train_loader, test_loader = get_data_loader(args) #feature_extraction = FeatureExtractionTest(train_loader, test_loader, args.cuda, args.batch_size) # Start model training if args.is_train == 'True': model.train(train_loader) # start evaluating on test data else: model.evaluate(test_loader, args.load_D, args.load_G)
def test(model: gan.GAN): model.eval() fake_data = model(64) plot_grid(fake_data.detach(), figsize=(8, 8), gridspec_kw=dict(wspace=.1, hspace=.1)) plt.show()
def main(args): #--------------prepare data------------------------ dataset = args.dataset dataroot = args.dataroot if not os.path.exists(dataroot): os.makedirs(dataroot) batch_size = args.batch_size epochs = args.epochs channels = args.channels model = None model_name = args.model if args.model == 'GAN': model = GAN(epochs, batch_size) elif args.model == 'DCGAN': model = DCGAN_MODEL(args) elif args.model == 'WGAN-CP': model = WGAN_CP(args) elif args.model == 'WGAN-GP': model = WGAN_GP(args) else: print("Model type non-existing. Try again.") exit(-1) workers = 0 # number of workers for dataloader, 2 creates problems utils = Utils() train_loader, test_loader = utils.prepare_data(dataroot, batch_size, workers, dataset, model_name, channels) # Start model training resume_training = False if args.resume_training == 'True': resume_training = True if args.is_train == 'True': model.train(train_loader, resume_training) # start evaluating on test data else: model.evaluate(test_loader, args.load_D, args.load_G)
def __init__(self, z_dim, crop_image_size, resized_image_size, batch_size, data_dir, clip_values=(-0.01, 0.01), critic_iterations=5, root_scope_name=''): self.clip_values = clip_values GAN.__init__(self, z_dim, crop_image_size, resized_image_size, batch_size, data_dir, critic_iterations, root_scope_name)
def make(self, cf, optimizer=None): if optimizer is None: raise ValueError('optimizer can not be None') if cf.model_name in [ 'lenet', 'alexNet', 'vgg16', 'vgg19', 'resnet50', 'InceptionV3', 'fcn8', 'unet', 'segnet_vgg', 'segnet_basic', 'resnetFCN', 'yolo', 'tiny-yolo' ]: in_shape, loss, metrics = self.basic_model_properties(cf, True) model = self.make_one_net_model(cf, in_shape, loss, metrics, optimizer) elif cf.model_name == 'adversarial_semseg': # loss, metrics and optimizer are made in class Adversarial_Semseg in_shape, _, _ = self.basic_model_properties(cf, False) model = Adversarial_Semseg(cf, in_shape) elif cf.model_name == 'gan': # loss, metrics and optimizer are made in class Adversarial_Semseg in_shape, _, _ = self.basic_model_properties(cf, False) model = GAN(cf, in_shape) else: raise ValueError('Unknown model name') # Output the model print(' Model: ' + cf.model_name) return model
def main(_): if FLAGS.network == 'vae': model = VAE(mode=FLAGS.mode, batch_size=FLAGS.batch_size, latent_dim=FLAGS.latent_dim) solver = VAE_Solver(model, batch_size=FLAGS.batch_size, train_iter=FLAGS.train_iter, log_dir=FLAGS.log_save_path, model_save_path=FLAGS.model_save_path, sample_save_path=FLAGS.sample_save_path) # create directories if not exist if not tf.gfile.Exists(FLAGS.model_save_path): tf.gfile.MakeDirs(FLAGS.model_save_path) if not tf.gfile.Exists(FLAGS.sample_save_path): tf.gfile.MakeDirs(FLAGS.sample_save_path) if FLAGS.mode == 'train': solver.train() elif FLAGS.mode == 'reconstruct': solver.reconstruct() elif FLAGS.mode == 'sample': solver.sample() elif FLAGS.mode == 'encode': solver.encode() elif FLAGS.network == 'gan': z_dim = 100 model = GAN(mode=FLAGS.mode) solver = GAN_Solver(model, batch_size=FLAGS.batch_size, z_dim=z_dim, train_iter=FLAGS.train_iter, log_dir=FLAGS.log_save_path, model_save_path=FLAGS.model_save_path, sample_save_path=FLAGS.sample_save_path) # create directories if not exist if not tf.gfile.Exists(FLAGS.model_save_path): tf.gfile.MakeDirs(FLAGS.model_save_path) if not tf.gfile.Exists(FLAGS.sample_save_path): tf.gfile.MakeDirs(FLAGS.sample_save_path) if FLAGS.mode == 'train': solver.train() elif FLAGS.mode == 'sample': solver.sample() elif FLAGS.network == 'acgan': z_dim = 128 feature_class = 'Smiling' model = ACGAN(mode=FLAGS.mode, batch_size=FLAGS.batch_size) solver = ACGAN_Solver(model, batch_size=FLAGS.batch_size, z_dim=z_dim, feature_class=feature_class, train_iter=FLAGS.train_iter, log_dir=FLAGS.log_save_path, model_save_path=FLAGS.model_save_path, sample_save_path=FLAGS.sample_save_path) # create directories if not exist if not tf.gfile.Exists(FLAGS.model_save_path): tf.gfile.MakeDirs(FLAGS.model_save_path) if not tf.gfile.Exists(FLAGS.sample_save_path): tf.gfile.MakeDirs(FLAGS.sample_save_path) if FLAGS.mode == 'train': solver.train() elif FLAGS.mode == 'sample': solver.sample()
def main(argv=None): gen_dim = FLAGS.gen_dimension generator_dims = [64 * gen_dim, 64 * gen_dim // 2, 64 * gen_dim // 4, 64 * gen_dim // 8, 3] discriminator_dims = [3, 64, 64 * 2, 64 * 4, 64 * 8, 1] crop_image_size, resized_image_size = map(int, FLAGS.image_size.split(',')) # if FLAGS.model == 0: # model = GAN(FLAGS.z_dim, crop_image_size, resized_image_size, FLAGS.batch_size, FLAGS.data_dir, critic_iterations=1) # elif FLAGS.model == 1: # model = WasserstienGAN(FLAGS.z_dim, crop_image_size, resized_image_size, FLAGS.batch_size, FLAGS.data_dir, # clip_values=(-0.01, 0.01), critic_iterations=25) # else: # raise ValueError("Unknown model identifier - FLAGS.model=%d" % FLAGS.model) # # model.create_network(generator_dims, discriminator_dims, FLAGS.optimizer, FLAGS.learning_rate, # FLAGS.optimizer_param) # # model.initialize_network(FLAGS.logs_dir, FLAGS.checkpoint_file) # # if FLAGS.mode == "train": # model.train_model(int(1 + FLAGS.iterations)) # elif FLAGS.mode == "visualize": # model.visualize_model() import cross_dis gan1_scope_name = 'gan' gan2_scope_name = 'wgan' with tf.variable_scope(gan1_scope_name): gan1 = GAN(FLAGS.z_dim, crop_image_size, resized_image_size, FLAGS.batch_size, FLAGS.data_dir, critic_iterations=1, root_scope_name='gan/') gan1.create_network(generator_dims, discriminator_dims, FLAGS.optimizer, FLAGS.learning_rate, FLAGS.optimizer_param) with tf.variable_scope(gan2_scope_name): gan2 = WasserstienGAN(FLAGS.z_dim, crop_image_size, resized_image_size, FLAGS.batch_size, FLAGS.data_dir, root_scope_name='wgan/', critic_iterations=5) gan2.create_network(generator_dims, discriminator_dims, FLAGS.optimizer, FLAGS.learning_rate, FLAGS.optimizer_param) cross_dis.run(gan1, gan2, gan1_scope_name, gan2_scope_name, discriminator_dims, FLAGS.logs_dir, FLAGS.checkpoint_file, int(FLAGS.iterations))
def __init__(self, args): GAN.__init__(self) self.build_model()
def train(cv_num_batch, tr_num_batch): with tf.Graph().as_default(): with tf.device('/cpu:0'): with tf.name_scope('input'): tr_data_list = read_list(FLAGS.tr_list_file) tr_inputs, tr_labels = get_batch( tr_data_list, FLAGS.batch_size, FLAGS.input_dim, FLAGS.output_dim, FLAGS.left_context, FLAGS.right_context, FLAGS.num_threads, FLAGS.max_epoches) cv_data_list = read_list(FLAGS.cv_list_file) cv_inputs, cv_labels = get_batch( cv_data_list, FLAGS.batch_size, FLAGS.input_dim, FLAGS.output_dim, FLAGS.left_context, FLAGS.right_context, FLAGS.num_threads, FLAGS.max_epoches) devices = [] for i in xrange(FLAGS.num_gpu): device_name = ("/gpu:%d" % i) print('Using device: ', device_name) devices.append(device_name) # Prevent exhausting all the gpu memories. config = tf.ConfigProto() config.gpu_options.allow_growth = True config.allow_soft_placement = True # execute the session with tf.Session(config=config) as sess: # Create two models with tr_inputs and cv_inputs individually. with tf.name_scope('model'): print( "=======================================================") print( "| Build Train model |") print( "=======================================================") tr_model = GAN(sess, FLAGS, devices, tr_inputs, tr_labels, cross_validation=False) # tr_model and val_model should share variables print( "=======================================================") print( "| Build Cross-Validation model |") print( "=======================================================") tf.get_variable_scope().reuse_variables() cv_model = GAN(sess, FLAGS, devices, cv_inputs, cv_labels, cross_validation=True) show_all_variables() init = tf.group(tf.global_variables_initializer(), tf.local_variables_initializer()) print("Initializing variables ...") sess.run(init) if tr_model.load(tr_model.save_dir): print("[*] Load SUCCESS") else: print("[!] Begin a new model.") coord = tf.train.Coordinator() threads = tf.train.start_queue_runners(sess=sess, coord=coord) try: cv_d_rl_loss, cv_d_fk_loss, \ cv_d_loss, cv_g_adv_loss, \ cv_g_mse_loss, cv_g_l2_loss, \ cv_g_loss = eval_one_epoch(sess, coord, cv_model, cv_num_batch, 0) print("CROSSVAL.LOSS PRERUN: " "d_rl_loss = {:.5f}, d_fk_loss = {:.5f}, " "d_loss = {:.5f}, g_adv_loss = {:.5f}, " "g_mse_loss = {:.5f}, g_l2_loss = {:.5f}, " "g_loss = {:.5f}".format(cv_d_rl_loss, cv_d_fk_loss, cv_d_loss, cv_g_adv_loss, cv_g_mse_loss, cv_g_l2_loss, cv_g_loss)) sys.stdout.flush() g_loss_prev = cv_g_loss decay_steps = 1 for epoch in range(FLAGS.max_epoches): start = datetime.datetime.now() tr_d_rl_loss, tr_d_fk_loss, \ tr_d_loss, tr_g_adv_loss, \ tr_g_mse_loss, tr_g_l2_loss, \ tr_g_loss = train_one_epoch(sess, coord, tr_model, tr_num_batch, epoch+1) cv_d_rl_loss, cv_d_fk_loss, \ cv_d_loss, cv_g_adv_loss, \ cv_g_mse_loss, cv_g_l2_loss, \ cv_g_loss = eval_one_epoch(sess, coord, cv_model, cv_num_batch, epoch+1) d_lr, g_lr = sess.run( [tr_model.d_learning_rate, tr_model.g_learning_rate]) end = datetime.datetime.now() print("Epoch {} (TRAIN AVG.LOSS): " "d_rl_loss = {:.5f}, d_fk_loss = {:.5f}, " "d_loss = {:.5f}, g_adv_loss = {:.5f}, " "g_mse_loss = {:.5f}, g_l2_loss = {:.5f}, " "g_loss = {:.5f}, " "d_lr = {:.3e}, g_lr = {:.3e}\n" "Epoch {} (CROSS AVG.LOSS): " "d_rl_loss = {:.5f}, d_fk_loss = {:.5f}, " "d_loss = {:.5f}, g_adv_loss = {:.5f}, " "g_mse_loss = {:.5f}, g_l2_loss = {:.5f}, " "g_loss = {:.5f}, " "time = {:.2f} h".format( epoch + 1, tr_d_rl_loss, tr_d_fk_loss, tr_d_loss, tr_g_adv_loss, tr_g_mse_loss, tr_g_l2_loss, tr_g_loss, d_lr, g_lr, epoch + 1, cv_d_rl_loss, cv_d_fk_loss, cv_d_loss, cv_g_adv_loss, cv_g_mse_loss, cv_g_l2_loss, cv_g_loss, (end - start).seconds / 3600.0)) sys.stdout.flush() g_loss_new = cv_g_loss # Accept or reject new parameters if g_loss_new < g_loss_prev: tr_model.save(tr_model.save_dir, epoch + 1) print("Epoch {}: Nnet Accepted. " "Save model SUCCESS.".format(epoch + 1)) # Relative loss between previous and current val_loss g_rel_impr = (g_loss_prev - g_loss_new) / g_loss_prev g_loss_prev = g_loss_new else: print("Epoch {}: Nnet Rejected.".format(epoch + 1)) if tr_model.load(tr_model.save_dir): print("[*] Load previous model SUCCESS.") sys.stdout.flush() else: print("[!] Load failed. No checkpoint from {} to " "restore previous model. Exit now.".format( tr_model.save_dir)) sys.stdout.flush() sys.exit(1) # Relative loss between previous and current val_loss g_rel_impr = (g_loss_prev - g_loss_new) / g_loss_prev # Start decay when improvement is low (Exponential decay) if g_rel_impr < FLAGS.start_decay_impr and \ epoch+1 >= FLAGS.keep_lr: g_learning_rate = \ FLAGS.g_learning_rate * \ FLAGS.decay_factor ** (decay_steps) d_learning_rate = \ FLAGS.d_learning_rate * \ FLAGS.decay_factor ** (decay_steps) disc_noise_std = \ FLAGS.init_disc_noise_std * \ FLAGS.decay_factor ** (decay_steps) sess.run( tf.assign(tr_model.g_learning_rate, g_learning_rate)) sess.run( tf.assign(tr_model.d_learning_rate, d_learning_rate)) sess.run( tf.assign(tr_model.disc_noise_std, disc_noise_std)) decay_steps += 1 # Stopping criterion if g_rel_impr < FLAGS.end_decay_impr: if epoch < FLAGS.min_epoches: print("Epoch %d: We were supposed to finish, " "but we continue as min_epoches %d" % (epoch + 1, FLAGS.min_epoches)) continue else: print("Epoch %d: Finished, too small relative " "G improvement %g" % (epoch + 1, g_rel_impr)) break except Exception, e: # Report exceptions to the coordinator. coord.request_stop(e) finally:
def decode(): """Decoding the inputs using current model.""" tf.logging.info("Get TEST sets number.") num_batch = get_num_batch(FLAGS.test_list_file, infer=True) with tf.Graph().as_default(): with tf.device('/cpu:0'): with tf.name_scope('input'): data_list = read_list(FLAGS.test_list_file) test_utt_id, test_inputs, _ = get_batch( data_list, batch_size=1, input_size=FLAGS.input_dim, output_size=FLAGS.output_dim, left=FLAGS.left_context, right=FLAGS.right_context, num_enqueuing_threads=FLAGS.num_threads, num_epochs=1, infer=True) devices = [] for i in xrange(FLAGS.num_gpu): device_name = ("/gpu:%d" % i) print('Using device: ', device_name) devices.append(device_name) # Prevent exhausting all the gpu memories. config = tf.ConfigProto() config.gpu_options.allow_growth = True config.allow_soft_placement = True # execute the session with tf.Session(config=config) as sess: # Create two models with tr_inputs and cv_inputs individually. with tf.name_scope('model'): model = GAN(sess, FLAGS, devices, test_inputs, labels=None, cross_validation=True) show_all_variables() init = tf.group(tf.global_variables_initializer(), tf.local_variables_initializer()) print("Initializing variables ...") sess.run(init) if model.load(model.save_dir, moving_average=True): print("[*] Load SUCCESS") else: print("[!] Load failed. Checkpoint not found. Exit now.") sys.exit(1) coord = tf.train.Coordinator() threads = tf.train.start_queue_runners(sess=sess, coord=coord) cmvn_filename = os.path.join(FLAGS.data_dir, "train_cmvn.npz") if os.path.isfile(cmvn_filename): cmvn = np.load(cmvn_filename) else: tf.logging.fatal("%s not exist, exit now." % cmvn_filename) sys.exit(1) out_dir_name = os.path.join(FLAGS.save_dir, 'test') if not os.path.exists(out_dir_name): os.makedirs(out_dir_name) write_scp_path = os.path.join(out_dir_name, 'feats.scp') write_ark_path = os.path.join(out_dir_name, 'feats.ark') writer = ArkWriter(write_scp_path) try: for batch in range(num_batch): if coord.should_stop(): break outputs = model.generator(test_inputs, None, reuse=True) outputs = tf.reshape(outputs, [-1, model.output_dim]) utt_id, activations = sess.run([test_utt_id, outputs]) sequence = activations * cmvn['stddev_labels'] + \ cmvn['mean_labels'] save_result = np.vstack(sequence) writer.write_next_utt(write_ark_path, utt_id[0], save_result) tf.logging.info("Write inferred %s to %s" % (utt_id[0], write_ark_path)) except Exception, e: # Report exceptions to the coordinator. coord.request_stop(e) finally:
def parser(): PARSER = argparse.ArgumentParser() # Training parameters PARSER.add_argument('--epochs', default=40, type=int, help='Number of training epochs.') PARSER.add_argument('--batch_size', type=int, default=64, help='batch size') PARSER.add_argument('--latent_dim', type=int, default=100, help='dimensionality of the latent space') PARSER.add_argument('--lr', type=float, default=0.0002, help='learning rate') PARSER.add_argument("--b1", type=float, default=0.5, help="momentum; beta1 in Adam optimizer.") PARSER.add_argument("--b2", type=float, default=0.999, help="decay; beta2 in Adam optimizer.") PARSER.add_argument('--dropout_D', type=float, default=0.2, help='Dropout probability on the Discriminator.') PARSER.add_argument('--dropout_G', type=float, default=0.2, help='Dropout probability on the Generator.') PARSER.add_argument('--label_smoothing', type=bool, default=True, help='Label Smoothing.') PARSER.add_argument('--flipped_labbels', type=bool, default=True, help='Flipped Labbels.') PARSER.add_argument('--eval_mode', type=bool, default=True, help='Evaluation mode On/Off when sampling.') PARSER.add_argument('--n_samples', type=int, default=9, help='The number of the generated images.') PARSER.add_argument('--device', default=None, type=str, help='Device to run the experiment. \ Valid options: "cpu", "cuda".') PARSER.add_argument('--seed', default=None, type=int, help='Fix random seed.') PARSER.add_argument('--model', default='gan', type=str, help="Model to be used. Valid options: \ 'gan'.") ARGS = PARSER.parse_args() if ARGS.device is None: ARGS.device = torch.device( "cuda:0" if torch.cuda.is_available() else "cpu") if ARGS.model == 'gan': from models.gan import GAN ARGS.model = nn.DataParallel(GAN(args=ARGS).to(ARGS.device)) else: print('Model {} is not implimented'.format(ARGS.model)) quit() print_(ARGS) return ARGS
def make_model(args, img_c): if args.model == 'autoencoder': if args.load: model = keras.models.load_model(os.path.join(args.out, 'ae')) print('Loaded model') else: # Autoencoder img = keras.Input((args.imsize, args.imsize, img_c), name='img-in') out = NormalizeImage()(img) out = encode(args, out, out_dim=args.zdim) out = MeasureNorm(name='latent_norm')(out) out = synthesize(args, out, img_c) out = AddMSE()((img, out)) model = keras.Model(img, out, name='autoencoder') model.compile(optimizer=keras.optimizers.Adam( args.ae_lr, args.beta1), steps_per_execution=args.steps_exec) print('Starting with new model') # Summarize model.summary() elif args.model == 'gan': # Generator and discriminator if args.load: # Loading model gen = keras.models.load_model(os.path.join(args.out, 'gen')) disc = keras.models.load_model(os.path.join(args.out, 'disc')) print('Loaded model') else: # Generator gen_in = keras.Input((args.imsize, args.imsize, img_c), name='gen-in') z = LatentMap(args)(gen_in) gen_out = synthesize(args, z, img_c) gen = keras.Model(gen_in, gen_out, name='generator') # Discriminator disc_in = keras.Input((args.imsize, args.imsize, img_c), name='disc-in') disc_out = NormalizeImage()(disc_in) disc_out = encode(args, disc_out, out_dim=1) disc = keras.Model(disc_in, disc_out, name='discriminator') gen.compile(keras.optimizers.Adam(args.gen_lr, args.beta1)) disc.compile(keras.optimizers.Adam(args.disc_lr, args.beta1)) print('Starting with new model') # Summarize disc.summary() gen.summary() # GAN model = GAN(args, gen, disc) model.compile(steps_per_execution=args.steps_exec) else: raise Exception(f'unknown model {args.model}') return model
def main(): parser = argparse.ArgumentParser(description="Generate 汉字 via generative adversarial network.") # Dataset parser.add_argument("--size", type=int, default=32, help="Font size.") parser.add_argument("--from_unicode", type=int, help="Starting point of the unicode.") parser.add_argument("--to_unicode", type=int, help="Ending point of the unicode.") parser.add_argument("--font", type=str, required=True, help="Path to the font file.") parser.add_argument("--num_workers", type=int, default=4, help="Number of data loading workers.") # Optimization parser.add_argument("--epochs", type=int, default=100, help="Number of epochs.") parser.add_argument("--batch_size", type=int, default=32, help="Batch size.") parser.add_argument("--gpu_ids", type=str, default='', help="GPUs for running this script.") parser.add_argument("--rand_dim", type=int, default=128, help="Dimension of the random vector.") parser.add_argument("--num_fakes", type=int, default=16, help="Use num_fakes generated images to train the discriminator.") parser.add_argument("--flip_rate", type=float, default=0.8, help="Label flipping rate.") parser.add_argument("--g_lr", type=float, default=0.01, help="Learning rate for generator.") parser.add_argument("--d_lr", type=float, default=0.01, help="Learning rate for discriminator.") parser.add_argument("--factor", type=float, default=0.2, help="Factor by which the learning rate will be reduced.") parser.add_argument("--patience", type=int, default=10, help="Number of epochs with no improvement after which learning rate will be reduced.") parser.add_argument("--threshold", type=float, default=0.1, help="Threshold for measuring the new optimum, to only focus on significant changes. ") # Misc parser.add_argument("--log_dir", type=str, default="../run/", help="Where to save the log?") parser.add_argument("--log_name", type=str, required=True, help="Name of the log folder.") parser.add_argument("--show_freq", type=int, default=64, help="How frequently to show generated images?") parser.add_argument("--seed", type=int, default=0, help="Random seed.") args = parser.parse_args() assert args.show_freq > 0 assert 0.0 <= args.flip_rate <= 1.0 # Check before run. if not os.path.exists(args.log_dir): os.mkdir(args.log_dir) log_dir = os.path.join(args.log_dir, args.log_name) # Setting up logger log_file = datetime.now().strftime("%Y-%m-%d-%H-%M-%S.log") sys.stdout = Logger(os.path.join(log_dir, log_file)) print(args) for s in args.gpu_ids: try: int(s) except ValueError as e: print("Invalid gpu id:{}".format(s)) raise ValueError os.environ['CUDA_VISIBLE_DEVICES'] = ','.join(args.gpu_ids) if args.gpu_ids: if torch.cuda.is_available(): use_gpu = True torch.cuda.manual_seed_all(args.seed) else: use_gpu = False else: use_gpu = False torch.manual_seed(args.seed) dataloader, size = build_dataloader(args.batch_size, args.num_workers, use_gpu, args.font, args.size, args.from_unicode, args.to_unicode) model = GAN(args.num_fakes, args.rand_dim, size, use_gpu) criterion = BCELoss() d_optimizer = torch.optim.SGD(model.discriminator.parameters(), lr=args.d_lr, momentum=0.9) g_optimizer = torch.optim.SGD(model.generator.parameters(), lr=args.g_lr, momentum=0.9) d_scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(d_optimizer, mode="min", factor=args.factor, patience=args.patience, verbose=True, threshold=args.threshold) g_scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(g_optimizer, mode="min", factor=args.factor, patience=args.patience, verbose=True, threshold=args.threshold) optimizer = d_optimizer, g_optimizer scheduler = d_scheduler, g_scheduler if use_gpu: model = model.cuda() model = torch.nn.DataParallel(model) print("Start training...") start = datetime.now() with SummaryWriter(log_dir) as writer: for epoch in range(args.epochs): for i, param_group in enumerate(d_optimizer.param_groups): d_learning_rate = float(param_group["lr"]) writer.add_scalar("d_lr_group_{0}".format(i), d_learning_rate, global_step=epoch) for i, param_group in enumerate(g_optimizer.param_groups): g_learning_rate = float(param_group["lr"]) writer.add_scalar("g_lr_group_{0}".format(i), g_learning_rate, global_step=epoch) train(model, dataloader, criterion, optimizer, use_gpu, writer, epoch, scheduler, args.num_fakes, args.flip_rate, args.show_freq) torch.save(model, os.path.join(log_dir, "latest.pth")) elapsed_time = str(datetime.now() - start) print("Finish training. Total elapsed time %s." % elapsed_time)
def test(model: gan.GAN): model.eval() fake_data = model(64) plot_grid(fake_data.detach(), figsize=(8, 8), gridspec_kw=dict(wspace=.1, hspace=.1)) plt.show()
def train(valdi_batch_per_iter, train_batch_per_iter, min_iters, max_iters): with tf.Graph().as_default(): with tf.device('/cpu:0'): with tf.name_scope('input'): tr_data_list = read_list(FLAGS.tr_list_file) tr_inputs, tr_labels = get_batch( tr_data_list, FLAGS.batch_size, FLAGS.input_dim, FLAGS.output_dim, FLAGS.left_context, FLAGS.right_context, FLAGS.num_threads, FLAGS.max_epoches) cv_data_list = read_list(FLAGS.cv_list_file) cv_inputs, cv_labels = get_batch( cv_data_list, FLAGS.batch_size, FLAGS.input_dim, FLAGS.output_dim, FLAGS.left_context, FLAGS.right_context, FLAGS.num_threads, None) devices = [] for i in xrange(FLAGS.num_gpu): device_name = ("/gpu:%d" % i) print('Using device: ', device_name) devices.append(device_name) # Prevent exhausting all the gpu memories. config = tf.ConfigProto() config.gpu_options.allow_growth = True config.allow_soft_placement = True # execute the session with tf.Session(config=config) as sess: # Create two models with tr_inputs and cv_inputs individually. with tf.name_scope('model'): print( "=======================================================") print( "| Build Train model |") print( "=======================================================") tr_model = GAN(sess, FLAGS, devices, tr_inputs, tr_labels, cross_validation=False) # tr_model and val_model should share variables print( "=======================================================") print( "| Build Cross-Validation model |") print( "=======================================================") tf.get_variable_scope().reuse_variables() cv_model = GAN(sess, FLAGS, devices, cv_inputs, cv_labels, cross_validation=True) show_all_variables() init = tf.group(tf.global_variables_initializer(), tf.local_variables_initializer()) print("Initializing variables ...") sess.run(init) if tr_model.load(tr_model.save_dir, moving_average=False): print("[*] Load SUCCESS") else: print("[!] Begin a new model.") coord = tf.train.Coordinator() threads = tf.train.start_queue_runners(sess=sess, coord=coord) try: # Early stop counter g_loss_prev = 10000.0 g_rel_impr = 1.0 check_interval = 3 windows_g_loss = [] g_learning_rate = FLAGS.num_gpu * FLAGS.g_learning_rate d_learning_rate = FLAGS.num_gpu * FLAGS.d_learning_rate sess.run(tf.assign(tr_model.g_learning_rate, g_learning_rate)) sess.run(tf.assign(tr_model.d_learning_rate, d_learning_rate)) for iteration in range(max_iters): start = datetime.datetime.now() tr_d_rl_loss, tr_d_fk_loss, \ tr_d_loss, tr_g_adv_loss, \ tr_g_mse_loss, tr_g_l2_loss, \ tr_g_loss = train_one_iteration(sess, coord, tr_model, train_batch_per_iter, iteration+1) cv_d_rl_loss, cv_d_fk_loss, \ cv_d_loss, cv_g_adv_loss, \ cv_g_mse_loss, cv_g_l2_loss, \ cv_g_loss = eval_one_iteration(sess, coord, cv_model, valdi_batch_per_iter, iteration+1) d_learning_rate, \ g_learning_rate = sess.run([tr_model.d_learning_rate, tr_model.g_learning_rate]) end = datetime.datetime.now() print("{}/{} (INFO): d_learning_rate = {:.5e}, " "g_learning_rate = {:.5e}, time = {:.3f} min\n" "{}/{} (TRAIN AVG.LOSS): " "d_rl_loss = {:.5f}, d_fk_loss = {:.5f}, " "d_loss = {:.5f}, g_adv_loss = {:.5f}, " "g_mse_loss = {:.5f}, g_l2_loss = {:.5f}, " "g_loss = {:.5f}\n" "{}/{} (CROSS AVG.LOSS): " "d_rl_loss = {:.5f}, d_fk_loss = {:.5f}, " "d_loss = {:.5f}, g_adv_loss = {:.5f}, " "g_mse_loss = {:.5f}, g_l2_loss = {:.5f}, " "g_loss = {:.5f}".format( iteration + 1, max_iters, d_learning_rate, g_learning_rate, (end - start).seconds / 60.0, iteration + 1, max_iters, tr_d_rl_loss, tr_d_fk_loss, tr_d_loss, tr_g_adv_loss, tr_g_mse_loss, tr_g_l2_loss, tr_g_loss, iteration + 1, max_iters, cv_d_rl_loss, cv_d_fk_loss, cv_d_loss, cv_g_adv_loss, cv_g_mse_loss, cv_g_l2_loss, cv_g_loss)) sys.stdout.flush() # Start decay learning rate g_learning_rate = exponential_decay( iteration + 1, FLAGS.num_gpu, min_iters, FLAGS.g_learning_rate) d_learning_rate = exponential_decay( iteration + 1, FLAGS.num_gpu, min_iters, FLAGS.d_learning_rate) disc_noise_std = exponential_decay( iteration + 1, FLAGS.num_gpu, min_iters, FLAGS.init_disc_noise_std, multiply_jobs=False) sess.run( tf.assign(tr_model.g_learning_rate, g_learning_rate)) sess.run( tf.assign(tr_model.d_learning_rate, d_learning_rate)) sess.run(tf.assign(tr_model.disc_noise_std, disc_noise_std)) windows_g_loss.append(cv_g_loss) # Accept or reject new parameters. if (iteration + 1) % check_interval == 0: g_loss_new = np.mean(windows_g_loss) g_rel_impr = (g_loss_prev - g_loss_new) / g_loss_prev if g_rel_impr > 0.0: tr_model.save(tr_model.save_dir, iteration + 1) print("Iteration {}: Nnet Accepted. " "Save model SUCCESS. g_loss_prev = {:.5f}, " "g_loss_new = {:.5f}".format( iteration + 1, g_loss_prev, g_loss_new)) g_loss_prev = g_loss_new else: print("Iteration {}: Nnet Rejected. " "g_loss_prev = {:.5f}, " "g_loss_new = {:.5f}".format( iteration + 1, g_loss_prev, g_loss_new)) # tr_model.load(tr_model.save_dir, moving_average=False) windows_g_loss = [] # Stopping criterion. if iteration + 1 > min_iters and \ (iteration + 1) % check_interval == 0: if g_rel_impr < FLAGS.end_improve: print("Iteration %d: Finished, too small relative " "G improvement %g" % (iteration + 1, g_rel_impr)) break sys.stdout.flush() if windows_g_loss: g_loss_new = np.mean(windows_g_loss) g_rel_impr = (g_loss_prev - g_loss_new) / g_loss_prev if g_rel_impr > 0.0: tr_model.save(tr_model.save_dir, iteration + 1) print("Iteration {}: Nnet Accepted. " "Save model SUCCESS. g_loss_prev = {:.5f}, " "g_loss_new = {:.5f}".format( iteration + 1, g_loss_prev, g_loss_new)) g_loss_prev = g_loss_new sys.stdout.flush() windows_g_loss = [] except Exception, e: # Report exceptions to the coordinator. coord.request_stop(e) finally:
ts = time.time() timestamp = datetime.datetime.fromtimestamp(ts).strftime('%d_%m_%Y_%H_%M_%S') parser = argparse.ArgumentParser(description='GAN without MI') parser.add_argument('--config', type=str, default='./configs/spiral_mine.yml', help = 'Path to config file') opts = parser.parse_args() params = get_config(opts.config) print(params) train_loader, val_loader = spiral_dataloader(params) if params['use_mine']: model = GAN_MI(params) else: model = GAN(params) if params['use_cuda']: model = model.cuda() logger = Logger(params['logs']) exp_logs = params['logs'] + params['exp_name'] + '_' + timestamp + '/' exp_results = params['results'] + params['exp_name'] + '_' + timestamp + '/' mkdir_p(exp_logs) mkdir_p(exp_results) if params['use_mine']: gan_trainer = GANTrainerMI(model, params, train_loader, val_loader, logger, exp_results, exp_logs) else: gan_trainer = GANTrainerVanilla(model, params, train_loader, val_loader, logger, exp_results, exp_logs)
def main(model_name): if not os.path.isdir(".ckpts"): os.mkdir(".ckpts") if model_name not in ["gan", "dcgan"]: print("The model name is wrong!") return ckpt_path = ".ckpts/%s/" % model_name if not os.path.isdir(ckpt_path): os.mkdir(ckpt_path) with open("config.json") as f: config = json.load(f)[model_name] loader = MNISTLoader() if model_name == "gan": model = GAN(loader.feature_depth, config["latent_depth"]) elif model_name == "dcgan": model = DCGAN(loader.feature_shape, config["latent_depth"]) steps_per_epoch = (loader.num_train_sets + loader.num_test_sets) // config["batch_size"] features = np.vstack([loader.train_features, loader.test_features]) features = feature_normalize(features) generator_losses_epoch = [] discriminator_losses_epoch = [] generated_images = [] for i in range(1, config["num_epochs"] + 1): generator_loss_epoch = [] discriminator_loss_epoch = [] for _ in range(steps_per_epoch): sampled_indices = \ np.random.choice( loader.num_train_sets + loader.num_test_sets, config["batch_size"], replace=False ) real_samples = features[sampled_indices] generator_loss, discriminator_loss = model.train_one_step( real_samples) generator_loss_epoch.append(generator_loss) discriminator_loss_epoch.append(discriminator_loss) generator_loss_epoch = np.mean(generator_loss_epoch) discriminator_loss_epoch = np.mean(discriminator_loss_epoch) print( "Epoch: %i, Generator Loss: %f, Discriminator Loss: %f" % \ (i, generator_loss_epoch, discriminator_loss_epoch) ) generator_losses_epoch.append(generator_loss_epoch) discriminator_losses_epoch.append(discriminator_loss_epoch) torch.save(model.generator.state_dict(), ckpt_path + "generator_%i.ckpt" % i) torch.save(model.discriminator.state_dict(), ckpt_path + "discriminator_%i.ckpt" % i) faked_samples = feature_denormalize( model.generate(config["batch_size"])) generated_images.append(faked_samples.detach().numpy()) with open(ckpt_path + "results.pkl", "wb") as f: pickle.dump((generator_losses_epoch, discriminator_losses_epoch, generated_images), f)
# crop_size = crop_size, # ratio = 0.5, # capacity = batch_size*10, # min_holding= batch_size*5, # threads = 8) dataset = MNISTDataSet('/home/nathan/envs/tensorflow/MNIST_data', batch_size = batch_size) network = GAN( sess = sess, zed_dim = 2, n_kernels = 64, bayesian = False, dataset = dataset, input_channel = 1, log_dir = log_dir, save_dir = save_dir, input_dims = [28,28], load_snapshot = False, learning_rate = 2e-4, label_dim = 10) ## Has to come after init_op ??? coord = tf.train.Coordinator() threads = tf.train.start_queue_runners(coord=coord) """ Training loop. Call network.train_step() once for each global step. Insert testing / snapshotting however you want.