def initialize_models(args, device): # network En_A = models.encoder(in_nc=args.in_ngc, nf=args.ngf, img_size=args.img_size).to(device) En_B = models.encoder(in_nc=args.in_ngc, nf=args.ngf, img_size=args.img_size).to(device) De_A = models.decoder(out_nc=args.out_ngc, nf=args.ngf).to(device) De_B = models.decoder(out_nc=args.out_ngc, nf=args.ngf).to(device) Disc_A = models.discriminator(in_nc=args.in_ndc, out_nc=args.out_ndc, nf=args.ndf, img_size=args.img_size).to(device) Disc_B = models.discriminator(in_nc=args.in_ndc, out_nc=args.out_ndc, nf=args.ndf, img_size=args.img_size).to(device) print('---------- models initialized -------------') utils.print_network(En_A) utils.print_network(En_B) utils.print_network(De_A) utils.print_network(De_B) utils.print_network(Disc_A) utils.print_network(Disc_B) print('-----------------------------------------------') # Parallelize code En_A = nn.DataParallel(En_A) En_B = nn.DataParallel(En_B) De_A = nn.DataParallel(De_A) De_B = nn.DataParallel(De_B) Disc_A = nn.DataParallel(Disc_A) Disc_B = nn.DataParallel(Disc_B) all_models = [En_A, En_B, De_A, De_B, Disc_A, Disc_B] return all_models
def get_t_d(conf, r_inputs, d_data): """ Create Real and fake temoral discriminators """ # to crop out unstable part for temporal discriminator, details in TecoGAN supplemental paper crop_size_dt = int(conf.train.crop_size * 4 * conf.gan.crop_dt) offset_dt = (conf.train.crop_size * 4 - crop_size_dt) // 2 crop_size_dt = conf.train.crop_size * 4 - offset_dt * 2 paddings = (0, 0, offset_dt, offset_dt, offset_dt, offset_dt, 0, 0) with nn.parameter_scope("discriminator"): real_warp = warp_by_flow(d_data.t_targets, d_data.t_vel) real_warp = space_to_depth_disc(real_warp, d_data.t_batch) # equivalent to tf.image.crop_to_bounding_box real_warp = real_warp[:, offset_dt:offset_dt + crop_size_dt, offset_dt:offset_dt + crop_size_dt, :] real_warp = F.pad(real_warp, paddings) before_warp = space_to_depth_disc(d_data.t_targets, d_data.t_batch) t_input = space_to_depth_disc(r_inputs[:, :d_data.t_size, :, :, :], d_data.t_batch) # resizing using bilinear interpolation input_hi = F.interpolate(t_input, scale=(4, 4), mode='linear', channel_last=True) real_warp = F.concatenate(before_warp, real_warp, input_hi) tdiscrim_real_output, real_layers = discriminator(real_warp) fake_warp = warp_by_flow(d_data.t_gen_output, d_data.t_vel) fake_warp = space_to_depth_disc(fake_warp, d_data.t_batch) fake_warp = fake_warp[:, offset_dt:offset_dt + crop_size_dt, offset_dt:offset_dt + crop_size_dt, :] fake_warp = F.pad(fake_warp, paddings) before_warp = space_to_depth_disc(d_data.t_gen_output, d_data.t_batch, inplace=False) fake_warp = F.concatenate(before_warp, fake_warp, input_hi) tdiscrim_fake_output, fake_layers = discriminator(fake_warp) temporal_disc = collections.namedtuple( 'temporal_disc', 'tdiscrim_real_output,' 'real_layers, tdiscrim_fake_output, fake_layers') return temporal_disc(tdiscrim_real_output=tdiscrim_real_output, real_layers=real_layers, tdiscrim_fake_output=tdiscrim_fake_output, fake_layers=fake_layers)
def train_step(images): noise = tf.random.normal([batch_size, noise_dim]) with tf.GradientTape() as gen_tape, tf.GradientTape() as disc_tape: generated_images = generator(noise, training=True) real_output = discriminator(images, training=True) fake_output = discriminator(generated_images, training=True) gen_loss = generator_loss(fake_output) disc_loss = discriminator_loss(real_output, fake_output) gradients_of_generator = gen_tape.gradient(gen_loss, generator.trainable_variables) gradients_of_discriminator = disc_tape.gradient(disc_loss, discriminator.trainable_variables) generator_optimizer.apply_gradients(zip(gradients_of_generator, generator.trainable_variables)) discriminator_optimizer.apply_gradients(zip(gradients_of_discriminator, discriminator.trainable_variables))
def train(epochs, batch_size, ckpt_path, imgs_path, lr, out_path): tf.keras.backend.clear_session() train_data = get_data(imgs_path, batch_size) gen = generator() disc = discriminator() print(gen.summary()) print(disc.summary()) gen_opt = Adam(learning_rate=lr, beta_1=0.5) disc_opt = Adam(learning_rate=lr, beta_1=0.5) ckpt = Checkpoint(disc=disc, gen=gen, disc_opt=disc_opt, gen_opt=gen_opt) manager = CheckpointManager(ckpt, ckpt_path, max_to_keep=3) if manager.latest_checkpoint: print("Restored from {}".format(manager.latest_checkpoint)) ckpt.restore(manager.latest_checkpoint) else: print("Initializing from scratch.") seed = tf.random.normal([16, ENCODING_SIZE], seed=1234) generate_and_save_images(gen, 0, seed, out_path) for ep in range(epochs): gen_loss = [] disc_loss_real = [] disc_loss_fake = [] print('Epoch: %d of %d' % (ep + 1, epochs)) start = time.time() for images in train_data: g_loss, d_loss_r, d_loss_f = train_step(images, gen, disc, gen_opt, disc_opt, batch_size) gen_loss.append(g_loss) disc_loss_real.append(d_loss_r) disc_loss_fake.append(d_loss_f) gen_loss = np.mean(np.asarray(gen_loss)) disc_loss_real = np.mean(np.asarray(disc_loss_real)) disc_loss_fake = np.mean(np.asarray(disc_loss_fake)) if (np.isnan(gen_loss) or np.isnan(disc_loss_real) or np.isnan(disc_loss_fake)): print("Something broke.") break manager.save() generate_and_save_images(gen, ep + 1, seed, out_path) print("Time for epoch:", time.time() - start) print("Gen loss=", gen_loss) print("Disc loss real=", disc_loss_real) print("Disc loss fake=", disc_loss_fake)
def __gradient_penalty_loss(self, real_data, fake_data, LAMBDA, i): real_data = tf.image.resize(real_data, size=(tf.shape(fake_data)[1], tf.shape(fake_data)[2])) alpha = tf.random.uniform(shape=(1, 1)) * tf.ones( shape=tf.shape(fake_data)) interpolates = alpha * real_data + (tf.ones_like(alpha) - alpha) * fake_data disc_interpolates = models.discriminator(interpolates, self.config, name="scale_" + str(i)) gradients = tf.gradients(disc_interpolates, interpolates)[0] return tf.reduce_mean(gradients**2) * LAMBDA
def train(args): train_ds, test_ds = get_data(args.img_path, args.batch) gen = generator() disc = discriminator() gen_opt = Adam(args.learning_rate, beta_1=0.5, beta_2=0.999) disc_opt = Adam(args.learning_rate, beta_1=0.5, beta_2=0.999) print(gen.summary()) print(disc.summary()) ckpt = Checkpoint(disc=disc, gen=gen, disc_opt=disc_opt, gen_opt=gen_opt) manager = CheckpointManager(ckpt, args.ckpt_path, max_to_keep=3) if args.continue_training: latest = manager.latest_checkpoint if latest: print("Restored from {}".format(latest)) ckpt.restore(latest) off = int(re.split('-', latest)[-1]) else: off = 0 print("Initializing from scratch.") for ep in range(args.epochs): for x, y in test_ds.take(1): generate_and_save_imgs(gen, ep + off, x, y, args.out_path) gen_loss = [] disc_loss = [] print('Epoch: %d of %d' % (ep + 1 + off, args.epochs + off)) start = time.time() for x, y in train_ds: g_loss, d_loss = train_step(x, y, gen, disc, gen_opt, disc_opt, args.batch) gen_loss.append(g_loss) disc_loss.append(d_loss) gen_loss = np.mean(np.asarray(gen_loss)) disc_loss = np.mean(np.asarray(disc_loss)) manager.save() print("Time for epoch:", time.time() - start) print("Gen loss=", gen_loss) print("Disc loss=", disc_loss) # Storing three different outputs after final epoch for x, y in test_ds.take(3): generate_and_save_imgs(gen, args.epochs + off, x, y, args.out_path) off += 1
def create_models(scene_size, modis_var_dim, noise_dim, lr_disc, lr_gan): # Create optimizers opt_disc = Adam(lr_disc, 0.5) opt_gan = Adam(lr_gan, 0.5) # Create models gen = models.cs_generator(scene_size, modis_var_dim, noise_dim) disc = models.discriminator(scene_size, modis_var_dim) # Compile models disc.trainable = False gan = models.cs_modis_cgan(gen, disc, scene_size, modis_var_dim, noise_dim) gan.compile(loss='binary_crossentropy', optimizer=opt_gan) disc.trainable = True disc.compile(loss='binary_crossentropy', optimizer=opt_disc) return (gen, disc, gan, opt_disc, opt_gan)
def train_wgan(input_shape, final_channels): generator = unet(input_shape, final_channels, use_pooling=False, skip_layers='inception', final_activation='tanh') disc_model = discriminator(INPUT_SHAPE, final_activation='linear') # generator.load_weights('degan/generator.h5') # discriminator.load_weights('degan/discriminator.h5') _model = DEGAN(generator, disc_model) gen_decay_rate = 5e-4 disc_decay_rate = 1e-4 _model.compile(tf.keras.optimizers.Adam(learning_rate=gen_decay_rate), tf.keras.optimizers.RMSprop(learning_rate=disc_decay_rate), wasserstein_gen_loss_fn, wasserstein_disc_loss_fn) gen_lr = CustomReduceLROnPlateau(_model.gen_optimizer, 'gen_lr', monitor='generator_loss', patience=200, factor=RATE_DECAY, verbose=1) disc_lr = CustomReduceLROnPlateau(_model.disc_optimizer, 'disc_lr', monitor='discriminator_loss', patience=200, factor=RATE_DECAY, verbose=1) model_path = 'degan/model_name.h5' model_checkpoints = ModelCheckpoint(model_path, monitor='generator_mae', save_best_only=True, verbose=1) tensorboard = TensorBoard(profile_batch='10, 20') callbacks = [gen_lr, disc_lr, tensorboard, model_checkpoints] return _model, callbacks
def gradient_penalty(self, real, generated): shape = tf.shape(generated)[0] if self.opts['aug_rate'] > 1.0: idxs = tf.range(shape) ridxs = tf.random_shuffle(idxs) real = tf.gather(real, ridxs) elif self.opts['aug_rate'] < 1.0: real_shape = tf.shape(real)[0] idxs = tf.range(real_shape) ridxs = tf.random_shuffle(idxs)[:shape] real = tf.gather(real, ridxs) alpha = tf.random_uniform(shape=[shape, 1, 1, 1], minval=0., maxval=1.) # alpha = tf.random_uniform(shape=[self.opts['batch_size'], 1, 1, 1], minval=0., maxval=1.) differences = generated - real interpolates = real + (alpha * differences) gradients = \ tf.gradients(discriminator(self.opts, interpolates, is_training=self.is_training)[1], [interpolates])[0] slopes = tf.sqrt(tf.reduce_sum(tf.square(gradients), reduction_indices=[1, 2, 3])) gradient_penalty = tf.reduce_mean((slopes - 1.) ** 2) return gradient_penalty
def generate(): BATCH_SIZE = 2 GENERATOR_PARAMS = 100 opt = Adam(lr=0.0002, beta_1=0.5, beta_2=0.999, epsilon=1e-08, decay=0.0) discriminator = models.discriminator() discriminator.compile(loss='binary_crossentropy', optimizer=opt) discriminator.load_weights('discriminator') generator = models.generator() generator.compile(loss='binary_crossentropy', optimizer=opt) generator.load_weights('generator') noise = np.zeros((BATCH_SIZE * 20, GENERATOR_PARAMS)) for i in range(BATCH_SIZE * 20): noise[i, :] = np.random.uniform(-1, 1, GENERATOR_PARAMS) generated_images = np.zeros((20 * BATCH_SIZE, 128, 128, 1)) for i in range(BATCH_SIZE * 20): generated_images[i, ...] = generator.predict(np.reshape( noise[i, :], (1, GENERATOR_PARAMS)), verbose=1) # d_pret = discriminator.predict(generated_images, verbose=1) # # index = np.arange(0, BATCH_SIZE * 20) # index.resize((BATCH_SIZE * 20, 1)) # pre_with_index = list(np.append(d_pret, index, axis=1)) # pre_with_index.sort(key=lambda x: x[0], reverse=True) # nice_images = np.zeros((BATCH_SIZE, 1) + (generated_images.shape[2:]), dtype=np.float32) for i in range(len(generated_images)): plt.imshow(generated_images[i, :, :, 0]) plt.axis('off') plt.savefig('E:/deep-neurons/generated-' + str(i) + '.png') plt.clf()
def model_fn(features, labels, mode, cfg): del labels resolution = features['resolution'] if mode == 'PREDICT': random_noise = features['random_noise'] * cfg.temperature return models.generator(random_noise, resolution, cfg, is_training=False) real_images_1 = features['real_images'] if cfg.data_format == 'NCHW': real_images_1 = utils.nchw_to_nhwc(real_images_1) real_images_2 = tf.image.flip_left_right(real_images_1) real_images_1 = utils.nhwc_to_nchw(real_images_1) real_images_2 = utils.nhwc_to_nchw(real_images_2) else: real_images_2 = tf.image.flip_left_right(real_images_1) random_noise_1 = features['random_noise_1'] fake_images_out_1 = models.generator(random_noise_1, resolution, cfg, is_training=True) real_scores_out = models.discriminator(real_images_1, resolution, cfg) fake_scores_out = models.discriminator(fake_images_out_1, resolution, cfg) #fake_scores_out_g = models.discriminator(fake_images_out_2, resolution, cfg) with tf.name_scope('Penalties'): d_loss = fake_scores_out - real_scores_out g_loss = -1.0 * fake_scores_out with tf.name_scope('GradientPenalty'): mixing_factors = tf.random_uniform( [int(real_images_1.get_shape()[0]), 1, 1, 1], 0.0, 1.0, dtype=fake_images_out_1.dtype) mixed_images_out = ops.lerp(real_images_1, real_images_2, mixing_factors) mixed_scores_out = models.discriminator(mixed_images_out, resolution, cfg) mixed_loss = tf.reduce_sum(mixed_scores_out) mixed_grads = tf.gradients(mixed_loss, [mixed_images_out])[0] mixed_norms = tf.sqrt( 1e-8 + tf.reduce_sum(tf.square(mixed_grads), axis=[1, 2, 3])) gradient_penalty = tf.square(mixed_norms - 1.0) d_loss += gradient_penalty * 10.0 with tf.name_scope('EpsilonPenalty'): epsilon_penalty = tf.square(real_scores_out) d_loss += epsilon_penalty * 0.001 resolution_step = utils.get_or_create_resolution_step() fadein_rate = tf.minimum( tf.cast(resolution_step, tf.float32) / float(cfg.fadein_steps), 1.0) learning_rate = cfg.base_learning_rate d_optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate, beta1=cfg.beta1, beta2=cfg.beta2, epsilon=cfg.eps, name="AdamD") g_optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate, beta1=cfg.beta1, beta2=cfg.beta2, epsilon=cfg.eps, name="AdamG") if cfg.data_format == 'NCHW': fake_images_out_1 = utils.nchw_to_nhwc(fake_images_out_1) real_images_1 = utils.nchw_to_nhwc(real_images_1) real_images_2 = utils.nchw_to_nhwc(real_images_2) mixed_images_out = utils.nchw_to_nhwc(mixed_images_out) tf.summary.image('generated_images', fake_images_out_1) tf.summary.image('real_images_1', real_images_1) tf.summary.image('real_images_2', real_images_2) tf.summary.image('mixed_images', mixed_images_out) with tf.variable_scope("Loss"): tf.summary.scalar('real_scores_out', tf.reduce_mean(real_scores_out)) tf.summary.scalar('fake_scores_out', tf.reduce_mean(fake_scores_out)) tf.summary.scalar('epsilon_penalty', tf.reduce_mean(epsilon_penalty)) tf.summary.scalar('mixed_norms', tf.reduce_mean(mixed_norms)) with tf.variable_scope("Rate"): tf.summary.scalar('fadein', fadein_rate) g_loss = tf.reduce_mean(g_loss) d_loss = tf.reduce_mean(d_loss) with tf.name_scope('TrainOps'): with tf.control_dependencies(tf.get_collection( tf.GraphKeys.UPDATE_OPS)): d_step = d_optimizer.minimize(d_loss, var_list=tf.get_collection( tf.GraphKeys.GLOBAL_VARIABLES, scope='Discriminator')) g_step = g_optimizer.minimize(g_loss, var_list=tf.get_collection( tf.GraphKeys.GLOBAL_VARIABLES, scope='Generator')) with tf.control_dependencies([g_step, d_step]): increment_global_step = tf.assign_add( tf.train.get_or_create_global_step(), 1) increment_resolution_step = tf.assign_add( utils.get_or_create_resolution_step(), 1) if resolution >= cfg.starting_resolution * 2: with tf.control_dependencies( [increment_global_step, increment_resolution_step]): lerp_ops = lerp_update_ops(resolution, fadein_rate) joint_op = tf.group([ d_step, g_step, lerp_ops[0], lerp_ops[1], increment_global_step, increment_resolution_step ]) else: joint_op = tf.group([ d_step, g_step, increment_global_step, increment_resolution_step ]) return joint_op, [g_loss, d_loss], [g_optimizer, d_optimizer]
from torch.autograd import Variable import itertools from models import encoder, decoder, discriminator, loss_functions import helpers import matplotlib.pyplot as plt import numpy as np from torch.nn import functional as F #%% setting up parameters batch_size, dim = 750, 2 Enc = encoder(dim=dim, k=2, batch_size=batch_size) Dec = decoder(dim=dim, k=2, batch_size=batch_size) Disc = discriminator(dim=dim, k=2, batch_size=batch_size) losses_ = loss_functions() dataHandler=helpers.data_and_plotting(batch_size,encoder=Enc,decoder=Dec,discriminator=Disc,mixture=False,\ semi_circle=True) #%% setting up the optimizers #generator optimizer optimizerE = optim.Adam(itertools.chain(Enc.parameters(), Dec.parameters()), lr=5e-4) #discriminator optimizer optimizerD = optim.Adam(itertools.chain(Disc.parameters()), lr=5e-4) ones = Variable(torch.FloatTensor(batch_size, 1).fill_(1.0), requires_grad=False)
def train_dcgan(n_epochs, batch_size, lr_rate, crop_len, scale_len, restore, paths): """ Train DCGAN. :param int n_epochs: Total number of epochs over the input data to train for. :param int batch_size: Batch size to use for training. :param float lr_rate: Generator learning rate. :param int crop_len: Image side length to use. :param int scale_len: Amount to scale the minimum side length to (for augmentation). :param bool restore: Specifies whether or not the latest checkpoint should be used. :param list paths: List of paths to images to use for training. """ assert scale_len >= crop_len, "invalid resize or crop length" # create placeholders sample = tf.placeholder(tf.float32, shape=[batch_size, Z_SIZE], name="sample") real = tf.placeholder(tf.float32, shape=[batch_size, 64, 64, 3], name="real") is_train = tf.placeholder(tf.bool, name="is_train") # instantiate the models G = generator(sample, is_train, crop_len) D_fake = discriminator(G, is_train) tf.get_variable_scope().reuse_variables() D_real = discriminator(real, is_train) # create losses loss_G = _sigmoid_loss(D_fake, tf.ones_like(D_fake)) loss_D = _sigmoid_loss(D_fake, tf.zeros_like(D_fake)) + \ _sigmoid_loss(D_real, tf.ones_like(D_real)) # acquire tensors for generator and discriminator # trick from carpedm20's implementation on github g_vars = [var for var in tf.trainable_variables() if "g_" in var.name] d_vars = [var for var in tf.trainable_variables() if "d_" in var.name] # create optimization objectives global_step = tf.Variable(0, name="global_step", trainable=False) opt_G = tf.train.AdamOptimizer(lr_rate, beta1=0.5).minimize(loss_G, var_list=g_vars) opt_D = tf.train.AdamOptimizer(2e-4, beta1=0.5).minimize(loss_D, var_list=d_vars) # create a saver and restore variables, if necessary saver = tf.train.Saver() model_path = os.path.join(OUTPUT_PATH, CHECKPOINT_NAME) # do some initialization sess = tf.Session() sess.run(tf.initialize_all_variables()) if restore: chkpt_fname = tf.train.latest_checkpoint(OUTPUT_PATH) saver.restore(sess, chkpt_fname) else: _clean_directory(OUTPUT_PATH) _clean_directory(OUTPUT_PATH) # reference vector (to examine epoch-to-epoch changes) vec_ref = np.random.normal(size=(batch_size, Z_SIZE)) # begin training n_iterations = len(paths) // batch_size for epoch in range(sess.run(global_step), n_epochs): print("------- EPOCH {0} -------".format(epoch)) random.shuffle(paths) # train the discriminator for one epoch for i in range(n_iterations): offset = (i * batch_size) % len(paths) batch_paths = paths[offset:(offset + batch_size)] imgs = _read_and_preprocess(batch_paths, scale_len, crop_len) vec = np.random.normal(size=(batch_size, Z_SIZE)) # minimize generator loss if i % TRAIN_RATIO == TRAIN_RATIO - 1: sess.run(opt_G, feed_dict={sample: vec, is_train: True}) # minimize discriminator loss sess.run(opt_D, feed_dict={ real: imgs, sample: vec, is_train: True }) # log the error if i % DISPLAY_LOSSES == 0: err_G = sess.run(loss_G, feed_dict={ sample: vec, is_train: False }) err_D = sess.run(loss_D, feed_dict={ real: imgs, sample: vec, is_train: False }) print(" Iteration {0}".format(i)) print(" generator loss = {0}".format(err_G)) print(" discriminator loss = {0}".format(err_D)) # save the model and sample results at the end of each epoch sess.run(tf.assign(global_step, epoch + 1)) saver.save(sess, model_path, global_step=global_step) batch_res = sess.run(G, {sample: vec_ref, is_train: False}) _deprocess_and_save(batch_res, epoch)
real_img = tf.placeholder(tf.float32, [None, 128, 128, 3], name='real_img') real_au = tf.placeholder(tf.float32, [None, 17], name='real_au') desired_au = tf.placeholder(tf.float32, [None, 17], name='desired_au') lr = tf.placeholder(tf.float32, name='lr') # ============= G & D ============= # G(Ic1, c2) * M fake_img, fake_mask = generator(real_img, desired_au, reuse=False) fake_img_masked = fake_mask * real_img + (1 - fake_mask) * fake_img # G(G(Ic1, c2)*M, c1) * M cyc_img, cyc_mask = generator(fake_img_masked, real_au, reuse=True) cyc_img_masked = cyc_mask * fake_img_masked + (1 - cyc_mask) * cyc_img # D(real_I) pred_real_img, pred_real_au = discriminator(real_img, reuse=False) # D(fake_I) pred_fake_img_masked, pred_fake_au = discriminator(fake_img_masked, reuse=True) # ============= losses ============= # G losses loss_g_fake_img_masked = -tf.reduce_mean(pred_fake_img_masked) * lambda_D_img loss_g_fake_au = l2_loss(desired_au, pred_fake_au) * lambda_D_au loss_g_cyc = l1_loss(real_img, cyc_img_masked) * lambda_cyc loss_g_mask_fake = tf.reduce_mean(fake_mask) * lambda_mask + smooth_loss(fake_mask) * lambda_mask_smooth loss_g_mask_cyc = tf.reduce_mean(cyc_mask) * lambda_mask + smooth_loss(cyc_mask) * lambda_mask_smooth loss_g = loss_g_fake_img_masked + loss_g_fake_au + \ loss_g_cyc + \
from models import generator, discriminator, GenNucleiDataset import torch G = generator() D = discriminator() G = torch.load('./app/weights/G_model.pth', map_location={'cuda:0': 'cpu'}) D = torch.load('./app/weights/D_model.pth', map_location={'cuda:0': 'cpu'}) print(G)
print "params:" print params dopt = Adam(lr=params.lr, beta_1=params.beta_1) # Define the U-Net generator unet = m.g_unet(params.a_ch, params.b_ch, params.nfatob, batch_size=params.batch_size, is_binary=params.is_b_binary) "unet summary:" unet.summary() # Define the discriminator d = m.discriminator(params.a_ch, params.b_ch, params.nfd, opt=dopt) "discriminator summary:" d.summary() if params.continue_train: load_weights(unet, d, log_dir=params.log_dir, expt_name=params.expt_name) ts = params.target_size train_dir = os.path.join(params.base_dir, params.train_dir) it_train = TwoImageIterator(train_dir, is_a_binary=params.is_a_binary, is_a_grayscale=params.is_a_grayscale, is_b_grayscale=params.is_b_grayscale,
def train(): graph = tf.Graph() with graph.as_default(): z = tf.placeholder(tf.float32, shape=[64, 100], name='z') img_batch = read_records.read_and_decode( 'tf_records/cartoon.tfrecords', batch_size=batch_size) #generator # fake=models.generator(z, stddev=0.02, alpha=alpha, name='generator', reuse=False) # # #discriminator # dis_real=models.discriminator(img_batch , alpha=alpha, batch_size=batch_size) # dis_fake=models.discriminator(fake, alpha=alpha, reuse=True) #generator fake = models.generator(z, reuse=False) #, is_training=True #discriminator dis_real = models.discriminator(img_batch, reuse=False) #is_training=True dis_fake = models.discriminator(fake, reuse=True) #, is_training=True # #losses # gene_loss = tf.reduce_mean(tf.squared_difference(dis_fake, 0.9)) # dis_loss = (tf.reduce_mean(tf.squared_difference(dis_real, 0.9)) # + tf.reduce_mean(tf.square(dis_fake))) / 2 gene_loss = tf.reduce_mean( tf.nn.sigmoid_cross_entropy_with_logits( labels=tf.ones_like(dis_fake) * 0.9, logits=dis_fake)) d_f_loss = tf.reduce_mean( tf.nn.sigmoid_cross_entropy_with_logits( labels=tf.zeros_like(dis_fake), logits=dis_fake)) d_r_loss = tf.reduce_mean( tf.nn.sigmoid_cross_entropy_with_logits( labels=tf.ones_like(dis_real) * 0.9, logits=dis_real)) dis_loss = d_f_loss + d_r_loss gen_loss_sum = tf.summary.scalar("gen_loss", gene_loss) dis_loss_sum = tf.summary.scalar("dis_loss", dis_loss) merge_sum_gen = tf.summary.merge([gen_loss_sum]) merge_sum_dis = tf.summary.merge([dis_loss_sum]) #variables gene_var = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope='generator') dis_var = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope='discriminator') gene_opt = tf.train.AdamOptimizer( learning_rate=0.0002, beta1=0.3).minimize(gene_loss, var_list=gene_var) dis_opt = tf.train.AdamOptimizer(learning_rate=0.0002, beta1=0.3).minimize(dis_loss, var_list=dis_var) test_sample = models.generator(z, reuse=True) #, is_training=False test_out = tf.add(test_sample, 0, 'test_out') init = tf.global_variables_initializer() print('t') with tf.Session(graph=graph) as sess: sess.run(init) # 初始化全局变量 z_ipt_sample = np.random.normal(size=[batch_size, 100]) coord = tf.train.Coordinator() threads = tf.train.start_queue_runners(coord=coord) writer = tf.summary.FileWriter('./tensorboard', sess.graph) saver = tf.train.Saver() try: for i in range(run_nums): z_ipt = np.random.normal(size=[batch_size, 100]) #train D #_, dis_loss1 = sess.run([dis_opt,dis_loss],feed_dict={real:img_batch,z:z_ipt}) sum_dis, _, dis_loss1 = sess.run( [merge_sum_dis, dis_opt, dis_loss], feed_dict={z: z_ipt}) #train G sum_gen, _, gen_loss1 = sess.run( [merge_sum_gen, gene_opt, gene_loss], feed_dict={z: z_ipt}) if i % 400 == 0: print(i) test_sample_opt = sess.run(test_sample, feed_dict={z: z_ipt_sample}) #print(type(test_sample_opt),test_sample_opt.shape) utils.mkdir('out_cartoon') utils.imwrite(utils.immerge(test_sample_opt, 10, 10), 'out_cartoon/' + str(i) + '.jpg') # writer.add_summary(sum_dis, i) #writer.add_summary(sum_gen, i) print("train end!!!") except tf.errors.OutOfRangeError: print('out of range') finally: coord.request_stop() coord.request_stop() coord.join(threads) writer.close() saver.save(sess, "./checkpoints/DCGAN")
parser.add_argument('--filename', type=str, default='monitor.npy.gz') args = parser.parse_args() # Defining ModelNet10 dataset for GAN dataset = ModelNet10GAN(filename=args.filename, dir=args.directory) # loading dataset into Dataloader data_loader = loader.DataLoader(dataset, batch_size=args.batch_size) num_epochs = args.epochs optimizerD = optim.Adam(D_.parameters(), lr=args.dis - lr, betas=(0.5, 0.999)) optimizerG = optim.Adam(G_.parameters(), lr=args.gen - lr, betas=(0.5, 0.999)) G_ = generator().to(device) D_ = discriminator().to(device) # Lists to store d_losses and g_losses. G_losses = [] D_losses = [] iters = 0 print("Starting training loop...") for epoch in range(num_epochs): # For each batch in epoch. RunningLossG = 0 RunningLossD = 0 for i, data in enumerate(data_loader, 1):
def main(): # Create output directory path_output = './checkpoints/' if not os.path.exists(path_output): os.makedirs(path_output) # Hyperparameters, to change epochs = 30 batch_size = 8 alpha = 1 # it's the trade-off parameter of loss function, what values should it take? gamma = 1 # Source domains name root = 'data/' source1 = 'real' source2 = 'sketch' source3 = 'infograph' target = 'quickdraw' # Dataloader dataset_s1 = dataset.DA(dir=root, name=source1, img_size=(224, 224), train=True) dataset_s2 = dataset.DA(dir=root, name=source2, img_size=(224, 224), train=True) dataset_s3 = dataset.DA(dir=root, name=source3, img_size=(224, 224), train=True) dataset_t = dataset.DA(dir=root, name=target, img_size=(224, 224), train=True) dataset_val = dataset.DA(dir=root, name=target, img_size=(224, 224), train=True, real_val=False) dataloader_s1 = DataLoader(dataset_s1, batch_size=batch_size, shuffle=True, num_workers=2) dataloader_s2 = DataLoader(dataset_s2, batch_size=batch_size, shuffle=True, num_workers=2) dataloader_s3 = DataLoader(dataset_s3, batch_size=batch_size, shuffle=True, num_workers=2) dataloader_t = DataLoader(dataset_t, batch_size=batch_size, shuffle=True, num_workers=2) dataloader_val = DataLoader(dataset_val, batch_size=batch_size, shuffle=False, num_workers=2) len_data = min(len(dataset_s1), len(dataset_s2), len(dataset_s3), len(dataset_t)) # length of "shorter" domain len_dataloader = min(len(dataloader_s1), len(dataloader_s2), len(dataloader_s3), len(dataloader_t)) # Define networks feature_extractor = models.feature_extractor() classifier_1 = models.class_classifier() classifier_2 = models.class_classifier() classifier_3 = models.class_classifier() classifier_1.apply(weight_init) classifier_2.apply(weight_init) classifier_3.apply(weight_init) discriminator = models.discriminator() discriminator.apply(weight_init) if torch.cuda.is_available(): feature_extractor = feature_extractor.cuda() classifier_1 = classifier_1.cuda() classifier_2 = classifier_2.cuda() classifier_3 = classifier_3.cuda() discriminator = discriminator.cuda() # Define loss cl_loss = nn.CrossEntropyLoss() disc_loss = nn.NLLLoss() # Optimizers # Change the LR optimizer_features = SGD(feature_extractor.parameters(), lr=0.0001, momentum=0.9, weight_decay=5e-4) optimizer_classifier = SGD(([{ 'params': classifier_1.parameters() }, { 'params': classifier_2.parameters() }, { 'params': classifier_3.parameters() }]), lr=0.002, momentum=0.9, weight_decay=5e-4) optimizer_discriminator = SGD(([ { 'params': discriminator.parameters() }, ]), lr=0.002, momentum=0.9, weight_decay=5e-4) # Lists training_loss = [] train_loss = [] train_class_loss = [] train_domain_loss = [] val_class_loss = [] val_domain_loss = [] acc_on_target = [] best_acc = 0.0 for epoch in range(epochs): epochTic = timeit.default_timer() tot_loss = 0.0 tot_c_loss = 0.0 tot_d_loss = 0.0 tot_val_c_loss = 0.0 tot_val_d_loss = 0.0 w1_mean = 0.0 w2_mean = 0.0 w3_mean = 0.0 feature_extractor.train() classifier_1.train(), classifier_2.train(), classifier_3.train() discriminator.train() if epoch + 1 == 5: optimizer_classifier = SGD(([{ 'params': classifier_1.parameters() }, { 'params': classifier_2.parameters() }, { 'params': classifier_3.parameters() }]), lr=0.001, momentum=0.9, weight_decay=5e-4) optimizer_discriminator = SGD( ([{ 'params': discriminator.parameters() }]), lr=0.001, momentum=0.9, weight_decay=5e-4) if epoch + 1 == 10: optimizer_classifier = SGD(([{ 'params': classifier_1.parameters() }, { 'params': classifier_2.parameters() }, { 'params': classifier_3.parameters() }]), lr=0.0001, momentum=0.9, weight_decay=5e-4) optimizer_discriminator = SGD( ([{ 'params': discriminator.parameters() }]), lr=0.0001, momentum=0.9, weight_decay=5e-4) print('*************************************************') for i, (data_1, data_2, data_3, data_t) in enumerate( zip(dataloader_s1, dataloader_s2, dataloader_s3, dataloader_t)): p = float(i + epoch * len_data) / epochs / len_data alpha = 2. / (1. + np.exp(-10 * p)) - 1 img1, lb1 = data_1 img2, lb2 = data_2 img3, lb3 = data_3 imgt, _ = data_t # Prepare data cur_batch = min(img1.shape[0], img2.shape[0], img3.shape[0], imgt.shape[0]) img1, lb1 = Variable(img1[0:cur_batch, :, :, :]).cuda(), Variable( lb1[0:cur_batch]).cuda() img2, lb2 = Variable(img2[0:cur_batch, :, :, :]).cuda(), Variable( lb2[0:cur_batch]).cuda() img3, lb3 = Variable(img3[0:cur_batch, :, :, :]).cuda(), Variable( lb3[0:cur_batch]).cuda() imgt = Variable(imgt[0:cur_batch, :, :, :]).cuda() # Forward optimizer_features.zero_grad() optimizer_classifier.zero_grad() optimizer_discriminator.zero_grad() # Extract Features ft1 = feature_extractor(img1) ft2 = feature_extractor(img2) ft3 = feature_extractor(img3) ft_t = feature_extractor(imgt) # Train the discriminator ds_s1 = discriminator(torch.cat((ft1, ft2, ft3)), alpha) ds_t = discriminator(ft_t, alpha) # Class Prediction cl1 = classifier_1(ft1) cl2 = classifier_2(ft2) cl3 = classifier_3(ft3) # Compute the "discriminator loss" ds_label = torch.zeros(cur_batch * 3).long() dt_label = torch.ones(cur_batch).long() d_s = disc_loss(ds_s1, ds_label.cuda()) d_t = disc_loss(ds_t, dt_label.cuda()) # Cross entropy loss l1 = cl_loss(cl1, lb1) l2 = cl_loss(cl2, lb2) l3 = cl_loss(cl3, lb3) # Classifier Weight total_class_loss = 1 / l1 + 1 / l2 + 1 / l3 w1 = (1 / l1) / total_class_loss w2 = (1 / l2) / total_class_loss w3 = (1 / l3) / total_class_loss w1_mean += w1 w2_mean += w2 w3_mean += w3 # total loss Class_loss = l1 + l2 + l3 Domain_loss = gamma * (d_s + d_t) loss = Class_loss + Domain_loss loss.backward() optimizer_features.step() optimizer_classifier.step() optimizer_discriminator.step() tot_loss += loss.item() * cur_batch tot_c_loss += Class_loss.item() tot_d_loss += Domain_loss.item() # Progress indicator print('\rTraining... Progress: %.1f %%' % (100 * (i + 1) / len_dataloader), end='') # Save Class loss and Domain loss if i % 50 == 0: train_class_loss.append(tot_c_loss / (i + 1)) train_domain_loss.append(tot_d_loss / (i + 1)) train_loss.append(tot_loss / (i + 1) / cur_batch) tot_t_loss = tot_loss / (len_data) training_loss.append(tot_t_loss) w1_mean /= len_dataloader w2_mean /= len_dataloader w3_mean /= len_dataloader #print(w1_mean,w2_mean,w3_mean) print('\rEpoch [%d/%d], Training loss: %.4f' % (epoch + 1, epochs, tot_t_loss), end='\n') #################################################################################################################### # Compute the accuracy at the end of each epoch feature_extractor.eval() classifier_1.eval(), classifier_2.eval(), classifier_3.eval() discriminator.eval() tot_acc = 0 with torch.no_grad(): for i, (imgt, lbt) in enumerate(dataloader_val): cur_batch = imgt.shape[0] imgt = imgt.cuda() lbt = lbt.cuda() # Forward the test images ft_t = feature_extractor(imgt) pred1 = classifier_1(ft_t) pred2 = classifier_2(ft_t) pred3 = classifier_3(ft_t) val_ds_t = discriminator(ft_t, alpha) # Compute class loss val_l1 = cl_loss(pred1, lbt) val_l2 = cl_loss(pred2, lbt) val_l3 = cl_loss(pred3, lbt) val_CE_loss = val_l1 + val_l2 + val_l3 # Compute domain loss val_dt_label = torch.ones(cur_batch).long() val_d_t = disc_loss(val_ds_t, val_dt_label.cuda()) # Compute accuracy output = pred1 * w1_mean + pred2 * w2_mean + pred3 * w3_mean _, pred = torch.max(output, dim=1) correct = pred.eq(lbt.data.view_as(pred)) accuracy = torch.mean(correct.type(torch.FloatTensor)) tot_acc += accuracy.item() * cur_batch # total loss tot_val_c_loss += val_CE_loss.item() tot_val_d_loss += val_d_t.item() # Progress indicator print('\rValidation... Progress: %.1f %%' % (100 * (i + 1) / len(dataloader_val)), end='') # Save validation loss if i % 50 == 0: val_class_loss.append(tot_val_c_loss / (i + 1)) val_domain_loss.append(tot_val_d_loss / (i + 1)) tot_t_acc = tot_acc / (len(dataset_val)) # Print acc_on_target.append(tot_t_acc) print('\rEpoch [%d/%d], Accuracy on target: %.4f' % (epoch + 1, epochs, tot_t_acc), end='\n') # Save every save_interval if best_acc < tot_t_acc: torch.save( { 'epoch': epoch, 'feature_extractor': feature_extractor.state_dict(), '{}_classifier'.format(source1): classifier_1.state_dict(), '{}_classifier'.format(source2): classifier_2.state_dict(), '{}_classifier'.format(source3): classifier_3.state_dict(), 'discriminator': discriminator.state_dict(), 'features_optimizer': optimizer_features.state_dict(), 'classifier_optimizer': optimizer_classifier.state_dict(), 'loss': training_loss, '{}_weight'.format(source1): w1_mean, '{}_weight'.format(source2): w2_mean, '{}_weight'.format(source3): w3_mean, }, os.path.join(path_output, target + '-{}-deming.pth'.format(epoch))) print('Saved best model!') best_acc = tot_t_acc # Pirnt elapsed time per epoch epochToc = timeit.default_timer() (t_min, t_sec) = divmod((epochToc - epochTic), 60) print('Elapsed time is: %d min: %d sec' % (t_min, t_sec)) # Save training loss and accuracy on target (if not 'real') pkl.dump(train_loss, open('{}total_loss_{}.p'.format(path_output, target), 'wb')) pkl.dump(train_class_loss, open('{}class_loss_{}.p'.format(path_output, target), 'wb')) pkl.dump(train_domain_loss, open('{}domain_loss_{}.p'.format(path_output, target), 'wb')) pkl.dump( acc_on_target, open('{}target_accuracy_{}.p'.format(path_output, target), 'wb')) pkl.dump( val_class_loss, open('{}val_class_loss_{}.p'.format(path_output, target), 'wb')) pkl.dump( val_domain_loss, open('{}val_domain_loss_{}.p'.format(path_output, target), 'wb'))
def train_dcgan(n_epochs, batch_size, lr_rate, crop_len, scale_len, restore, paths): """ Train DCGAN. :param int n_epochs: Total number of epochs over the input data to train for. :param int batch_size: Batch size to use for training. :param float lr_rate: Generator learning rate. :param int crop_len: Image side length to use. :param int scale_len: Amount to scale the minimum side length to (for augmentation). :param bool restore: Specifies whether or not the latest checkpoint should be used. :param list paths: List of paths to images to use for training. """ assert scale_len >= crop_len, "invalid resize or crop length" # create placeholders sample = tf.placeholder(tf.float32, shape=[batch_size, Z_SIZE], name="sample") real = tf.placeholder(tf.float32, shape=[batch_size, 64, 64, 3], name="real") is_train = tf.placeholder(tf.bool, name="is_train") # instantiate the models G = generator(sample, is_train, crop_len) D_fake = discriminator(G, is_train) tf.get_variable_scope().reuse_variables() D_real = discriminator(real, is_train) # create losses loss_G = _sigmoid_loss(D_fake, tf.ones_like(D_fake)) loss_D = _sigmoid_loss(D_fake, tf.zeros_like(D_fake)) + \ _sigmoid_loss(D_real, tf.ones_like(D_real)) # acquire tensors for generator and discriminator # trick from carpedm20's implementation on github g_vars = [var for var in tf.trainable_variables() if "g_" in var.name] d_vars = [var for var in tf.trainable_variables() if "d_" in var.name] # create optimization objectives global_step = tf.Variable(0, name="global_step", trainable=False) opt_G = tf.train.AdamOptimizer(lr_rate, beta1=0.5).minimize(loss_G, var_list=g_vars) opt_D = tf.train.AdamOptimizer(2e-4, beta1=0.5).minimize(loss_D, var_list=d_vars) # create a saver and restore variables, if necessary saver = tf.train.Saver() model_path = os.path.join(OUTPUT_PATH, CHECKPOINT_NAME) # do some initialization sess = tf.Session() sess.run(tf.initialize_all_variables()) if restore: chkpt_fname = tf.train.latest_checkpoint(OUTPUT_PATH) saver.restore(sess, chkpt_fname) else: _clean_directory(OUTPUT_PATH) _clean_directory(OUTPUT_PATH) # reference vector (to examine epoch-to-epoch changes) vec_ref = np.random.normal(size=(batch_size, Z_SIZE)) # begin training n_iterations = len(paths) // batch_size for epoch in range(sess.run(global_step), n_epochs): print("------- EPOCH {0} -------".format(epoch)) random.shuffle(paths) # train the discriminator for one epoch for i in range(n_iterations): offset = (i * batch_size) % len(paths) batch_paths = paths[offset:(offset + batch_size)] imgs = _read_and_preprocess(batch_paths, scale_len, crop_len) vec = np.random.normal(size=(batch_size, Z_SIZE)) # minimize generator loss if i % TRAIN_RATIO == TRAIN_RATIO - 1: sess.run(opt_G, feed_dict={sample: vec, is_train: True}) # minimize discriminator loss sess.run(opt_D, feed_dict={real: imgs, sample: vec, is_train: True}) # log the error if i % DISPLAY_LOSSES == 0: err_G = sess.run(loss_G, feed_dict={sample: vec, is_train: False}) err_D = sess.run(loss_D, feed_dict={real: imgs, sample: vec, is_train: False}) print(" Iteration {0}".format(i)) print(" generator loss = {0}".format(err_G)) print(" discriminator loss = {0}".format(err_D)) # save the model and sample results at the end of each epoch sess.run(tf.assign(global_step, epoch + 1)) saver.save(sess, model_path, global_step=global_step) batch_res = sess.run(G, {sample: vec_ref, is_train: False}) _deprocess_and_save(batch_res, epoch)
def train(args): # Communicator and Context extension_module = "cudnn" ctx = get_extension_context(extension_module, type_config=args.type_config) comm = C.MultiProcessDataParalellCommunicator(ctx) comm.init() n_devices = comm.size mpi_rank = comm.rank device_id = comm.local_rank ctx.device_id = str(device_id) nn.set_default_context(ctx) # Args latent = args.latent maps = args.maps batch_size = args.batch_size image_size = args.image_size n_classes = args.n_classes not_sn = args.not_sn # Model # workaround to start with the same weights in the distributed system. np.random.seed(412) # generator loss z = nn.Variable([batch_size, latent]) y_fake = nn.Variable([batch_size]) x_fake = generator(z, y_fake, maps=maps, n_classes=n_classes, sn=not_sn).apply(persistent=True) p_fake = discriminator(x_fake, y_fake, maps=maps // 16, n_classes=n_classes, sn=not_sn) loss_gen = gan_loss(p_fake) # discriminator loss y_real = nn.Variable([batch_size]) x_real = nn.Variable([batch_size, 3, image_size, image_size]) p_real = discriminator(x_real, y_real, maps=maps // 16, n_classes=n_classes, sn=not_sn) loss_dis = gan_loss(p_fake, p_real) # generator with fixed value for test z_test = nn.Variable.from_numpy_array(np.random.randn(batch_size, latent)) y_test = nn.Variable.from_numpy_array( generate_random_class(n_classes, batch_size)) x_test = generator(z_test, y_test, maps=maps, n_classes=n_classes, test=True, sn=not_sn) # Solver solver_gen = S.Adam(args.lrg, args.beta1, args.beta2) solver_dis = S.Adam(args.lrd, args.beta1, args.beta2) with nn.parameter_scope("generator"): params_gen = nn.get_parameters() solver_gen.set_parameters(params_gen) with nn.parameter_scope("discriminator"): params_dis = nn.get_parameters() solver_dis.set_parameters(params_dis) # Monitor if comm.rank == 0: monitor = Monitor(args.monitor_path) monitor_loss_gen = MonitorSeries( "Generator Loss", monitor, interval=10) monitor_loss_dis = MonitorSeries( "Discriminator Loss", monitor, interval=10) monitor_time = MonitorTimeElapsed( "Training Time", monitor, interval=10) monitor_image_tile_train = MonitorImageTile("Image Tile Train", monitor, num_images=args.batch_size, interval=1, normalize_method=normalize_method) monitor_image_tile_test = MonitorImageTile("Image Tile Test", monitor, num_images=args.batch_size, interval=1, normalize_method=normalize_method) # DataIterator rng = np.random.RandomState(device_id) di = data_iterator_imagenet(args.train_dir, args.dirname_to_label_path, args.batch_size, n_classes=args.n_classes, rng=rng) # Train loop for i in range(args.max_iter): # Train discriminator x_fake.need_grad = False # no need for discriminator backward solver_dis.zero_grad() for _ in range(args.accum_grad): # feed x_real and y_real x_data, y_data = di.next() x_real.d, y_real.d = x_data, y_data.flatten() # feed z and y_fake z_data = np.random.randn(args.batch_size, args.latent) y_data = generate_random_class(args.n_classes, args.batch_size) z.d, y_fake.d = z_data, y_data loss_dis.forward(clear_no_need_grad=True) loss_dis.backward( 1.0 / (args.accum_grad * n_devices), clear_buffer=True) comm.all_reduce([v.grad for v in params_dis.values()]) solver_dis.update() # Train genrator x_fake.need_grad = True # need for generator backward solver_gen.zero_grad() for _ in range(args.accum_grad): z_data = np.random.randn(args.batch_size, args.latent) y_data = generate_random_class(args.n_classes, args.batch_size) z.d, y_fake.d = z_data, y_data loss_gen.forward(clear_no_need_grad=True) loss_gen.backward( 1.0 / (args.accum_grad * n_devices), clear_buffer=True) comm.all_reduce([v.grad for v in params_gen.values()]) solver_gen.update() # Synchronize by averaging the weights over devices using allreduce if i % args.sync_weight_every_itr == 0: weights = [v.data for v in nn.get_parameters().values()] comm.all_reduce(weights, division=True, inplace=True) # Save model and image if i % args.save_interval == 0 and comm.rank == 0: x_test.forward(clear_buffer=True) nn.save_parameters(os.path.join( args.monitor_path, "params_{}.h5".format(i))) monitor_image_tile_train.add(i, x_fake.d) monitor_image_tile_test.add(i, x_test.d) # Monitor if comm.rank == 0: monitor_loss_gen.add(i, loss_gen.d.copy()) monitor_loss_dis.add(i, loss_dis.d.copy()) monitor_time.add(i) if comm.rank == 0: x_test.forward(clear_buffer=True) nn.save_parameters(os.path.join( args.monitor_path, "params_{}.h5".format(i))) monitor_image_tile_train.add(i, x_fake.d) monitor_image_tile_test.add(i, x_test.d)
height=flow_height, width=flow_width, reuse=None) train_pred_flow = flownet(input_a=train_inputs[..., -3:], input_b=train_outputs, height=flow_height, width=flow_width, reuse=True) flow_loss = tf.reduce_mean(tf.abs(train_gt_flow - train_pred_flow)) else: flow_loss = tf.constant(0.0, dtype=tf.float32) # define adversarial loss if adversarial: with tf.variable_scope('discriminator', reuse=None): real_logits, real_outputs = discriminator(inputs=train_gt) with tf.variable_scope('discriminator', reuse=True): fake_logits, fake_outputs = discriminator(inputs=train_outputs) print('real_outputs = {}'.format(real_outputs)) print('fake_outputs = {}'.format(fake_outputs)) adv_loss = tf.reduce_mean(tf.square(fake_outputs - 1) / 2) dis_loss = tf.reduce_mean(tf.square(real_outputs - 1) / 2) + tf.reduce_mean(tf.square(fake_outputs) / 2) else: adv_loss = tf.constant(0.0, dtype=tf.float32) dis_loss = tf.constant(0.0, dtype=tf.float32) with tf.name_scope('training'): g_loss = tf.add_n([
with tf.device('/gpu:%d' % gpu_id): ''' graph ''' # nodes a_real = tf.placeholder(tf.float32, shape=[None, crop_size, crop_size, 3]) b_real = tf.placeholder(tf.float32, shape=[None, crop_size, crop_size, 3]) a2b_sample = tf.placeholder(tf.float32, shape=[None, crop_size, crop_size, 3]) b2a_sample = tf.placeholder(tf.float32, shape=[None, crop_size, crop_size, 3]) a2b = models.generator(a_real, 'a2b') b2a = models.generator(b_real, 'b2a') b2a2b = models.generator(b2a, 'a2b', reuse=True) a2b2a = models.generator(a2b, 'b2a', reuse=True) a_dis = models.discriminator(a_real, 'a') b2a_dis = models.discriminator(b2a, 'a', reuse=True) b2a_sample_dis = models.discriminator(b2a_sample, 'a', reuse=True) b_dis = models.discriminator(b_real, 'b') a2b_dis = models.discriminator(a2b, 'b', reuse=True) a2b_sample_dis = models.discriminator(a2b_sample, 'b', reuse=True) # losses g_loss_a2b = tf.identity(ops.l2_loss(a2b_dis, tf.ones_like(a2b_dis)), name='g_loss_a2b') g_loss_b2a = tf.identity(ops.l2_loss(b2a_dis, tf.ones_like(b2a_dis)), name='g_loss_b2a') cyc_loss_a = tf.identity(ops.l1_loss(a_real, a2b2a) * 10.0, name='cyc_loss_a') cyc_loss_b = tf.identity(ops.l1_loss(b_real, b2a2b) * 10.0, name='cyc_loss_b')
import os os.environ['CUDA_VISIBLE_DEVICES']='0' import tensorflow as tf import numpy as np import cv2 import face_recognition from models import generator, discriminator from tqdm import tqdm real_img = tf.placeholder(tf.float32, [1, 128, 128, 3], name='real_img') style_img = tf.placeholder(tf.float32, [1, 128, 128, 3], name='style_img') _, desired_au = discriminator(style_img, reuse=False) fake_img, fake_mask = generator(real_img, desired_au, reuse=False) fake_img_masked = fake_mask * real_img + (1 - fake_mask) * fake_img imgs_names = os.listdir('jiaqi_face') real_src = face_recognition.load_image_file('obama.jpeg') # RGB image face_loc = face_recognition.face_locations(real_src) if len(face_loc) == 1: top, right, bottom, left = face_loc[0] real_face = np.zeros((1,128,128,3), dtype=np.float32) style_face = np.zeros((1,128,128,3), dtype=np.float32) real_face[0] = cv2.resize(real_src[top:bottom, left:right], (128,128)) / 127.5 - 1 saver = tf.train.Saver() with tf.Session() as sess: saver.restore(sess, 'weights/ganimation_upsample_epoch39.ckpt')
model = load_model(model_path) else: GAN = False encoder = False training_steps = 350 generator_model = None adversarial_model = None discriminator_model = None batch_size = 128 if GAN: for i in range(training_steps): x = X[i * batch_size:(i + 1) * batch_size, :, :] size = x.shape[0] if not discriminator_model: discriminator_model = discriminator() if not generator_model: generator_model = generator() if not adversarial_model: adversarial_model = adversarial( generator_model, discriminator_model) noise = np.random.uniform(0, 1.0, size=[batch_size, 100]) fakes = generator_model.predict(noise) x = np.concatenate((x, fakes)) x = np.expand_dims(x, axis=3) y = np.zeros([2 * batch_size, 1]) y[:batch_size, :] = 1 d_loss = discriminator_model.train_on_batch(x, y)
generator = eval('models.generator_%d'%opt.image_size) input_tensor_A = tf.placeholder(tf.float32, [None, opt.image_size, opt.image_size, 3]) input_tensor_B = tf.placeholder(tf.float32, [None, opt.image_size, opt.image_size, 3]) fake_pool_A = tf.placeholder(tf.float32, [None, opt.image_size, opt.image_size, 3]) fake_pool_B = tf.placeholder(tf.float32, [None, opt.image_size, opt.image_size, 3]) fake_B = generator(input_tensor_A, 'generator_B', skip=opt.skip) fake_A = generator(input_tensor_B, 'generator_A', skip=opt.skip) cycle_A = generator(fake_B, 'generator_A', reuse=True, skip=opt.skip) cycle_B = generator(fake_A, 'generator_B', reuse=True, skip=opt.skip) # define discriminator_A's loss fake_pool_A_dis = models.discriminator(fake_pool_A, 'discriminator_A') fake_A_dis = models.discriminator(fake_A, 'discriminator_A', reuse=True) real_A_dis = models.discriminator(input_tensor_A, 'discriminator_A', reuse=True) A_dis_loss_real = tf.reduce_mean(tf.square(real_A_dis - 1)) A_dis_loss_fake_pool = tf.reduce_mean(tf.square(fake_pool_A_dis)) A_dis_loss = (A_dis_loss_fake_pool + A_dis_loss_real) * 0.5 # define discriminator_B's loss fake_pool_B_dis = models.discriminator(fake_pool_B, 'discriminator_B') fake_B_dis = models.discriminator(fake_B, 'discriminator_B', reuse=True) real_B_dis = models.discriminator(input_tensor_B, 'discriminator_B', reuse=True) B_dis_loss_real = tf.reduce_mean(tf.square(real_B_dis - 1)) B_dis_loss_fake_pool = tf.reduce_mean(tf.square(fake_pool_B_dis)) B_dis_loss = (B_dis_loss_fake_pool + B_dis_loss_real) * 0.5
def __init__(self, num_latent, num_out, batch_size, num_disc, num_channels=3, num_hidden=1024, D_weights=None, G_weights=None, name='GMAN', mixing='arithmetic', weight_type='normal', objective='original', boosting_variant=None, self_challenged=False): self.num_latent = num_latent self.side = num_out self.num_channels = num_channels self.num_hidden = num_hidden self.batch_size = batch_size self.N = num_disc self.base_prob = 0.4 self.delta_p = (0.6 - self.base_prob) / self.N self.h_adv = num_hidden self.name = name self.weight_type = weight_type self.channel_size = self.batch_size * self.N self.self_challenged = self_challenged # boosting variables self.aux_vars = [] self.aux_vars_new = [] with tf.variable_scope(self.name): # Define latent distribution self.z = tf.random_uniform(shape=[self.channel_size, 1, 1, self.num_latent], minval=-1., maxval=1., name='z') # Generate fake images self.fake = generator(self) # print(f'self.fake -> {self.fake}') if boosting_variant is None: fake_split = tf.split(axis=0, num_or_size_splits=self.N, value=self.fake, name='fake_split') else: fake_split = [self.fake] * self.N # Discriminate fake images self.Df_logits = [discriminator(self, fake_split[ind], ind, (self.base_prob + self.delta_p * (ind + 1)), reuse=False) for ind in range(self.N)] # Retrieve real images self.real = tf.placeholder(tf.float32, shape=[self.channel_size, self.side, self.side, self.num_channels], name='real') if boosting_variant is None: real_split = tf.split(axis=0, num_or_size_splits=self.N, value=self.real, name='real_split') else: real_split = [self.real] * self.N # Discriminate real images self.Dr_logits = [discriminator(self, real_split[ind], ind, (self.base_prob + self.delta_p * (ind + 1)), reuse=True) for ind in range(self.N)] # Retrieve trainable weights t_vars = tf.trainable_variables() self.G_vars = [var for var in t_vars if (self.name + '/generator') in var.name] self.D_vars = [[var for var in t_vars if ('%s/discriminator_%d' % (self.name, num)) in var.name] for num in range(self.N)] # Assign values to weights (if given) self.assign_weights = [] if D_weights is not None: for i in range(self.N): for j in range(len(self.D_vars[i])): self.assign_weights.append(tf.assign(self.D_vars[i][j], D_weights[i][j])) if G_weights is not None: for j in range(len(self.G_vars)): self.assign_weights.append(tf.assign(self.G_vars[j], G_weights[j])) # Define Discriminator losses with tf.variable_scope('D_Loss'): if boosting_variant is None: self.get_D_losses(obj=objective) else: self.get_D_boosted_losses(boosting_variant, obj=objective) # Define Generator losses with tf.variable_scope('G_Loss'): if boosting_variant is None: self.get_G_loss(mixing, obj=objective) else: self.get_G_boosted_loss(boosting_variant, mixing, obj=objective) # Add Summaries self.add_summaries() # Construct Discriminator updates self.lrd = tf.placeholder(dtype=tf.float32) self.D_optimizer = tf.train.AdamOptimizer(learning_rate=self.lrd, beta1=0.5) self.D_optim = [self.D_optimizer.minimize(D_loss, var_list=self.D_vars[ind]) for ind, D_loss in enumerate(self.D_losses)] # Construct Generator updates self.lrg = tf.placeholder(dtype=tf.float32) self.G_optimizer = tf.train.AdamOptimizer(learning_rate=self.lrg, beta1=0.5) self.G_optim = self.G_optimizer.minimize(self.G_loss, var_list=self.G_vars) self.G_grads = self.G_optimizer.compute_gradients(self.G_loss, var_list=self.G_vars) # Join all updates self.all_opt = [opt for opt in self.D_optim] self.all_opt.append(self.G_optim)
gen_outputs = torch.stack(gen_outputs, dim=1) gen_outputs = gen_outputs.cpu().squeeze(0) save_as_gif(gen_outputs, f"./output/output{batch_idx}{args.videotype}") # My training loop for TecoGan elif args.mode == "train": # Defining dataset and dataloader dataset = train_dataset(args) dataloader = DataLoader(dataset, batch_size=4, shuffle=True, num_workers=8) # Defining the models as well as the optimizers and lr schedulers generator_F = generator(3, args=args).cuda() # fnet = f_net().cuda() discriminator_F = discriminator(args=args).cuda() counter1 = 0. counter2 = 0. min_gen_loss = np.inf tdis_learning_rate = args.learning_rate if not args.Dt_mergeDs: tdis_learning_rate = tdis_learning_rate * 0.3 tdiscrim_optimizer = torch.optim.Adam(discriminator_F.parameters(), tdis_learning_rate, betas=(args.beta, 0.999), eps=args.adameps) gen_optimizer = torch.optim.Adam(generator_F.parameters(), args.learning_rate, betas=(args.beta, 0.999), eps=args.adameps) # fnet_optimizer = torch.optim.Adam(fnet.parameters(), args.learning_rate, betas=(args.beta, 0.999),
def main(DATASET, data_size, num_train_epochs, num_disc_steps, num_gen_steps, batch_size, save_checkpoint_every, generate_samples_every, flip_alpha, restore = False): loader = load_Data_new.DataLoader(DATASET) # Create Placeholders for Data X = tf.placeholder(tf.float32, shape = [None, 64, 64, 3]) Z = tf.placeholder(tf.float32, shape = [None, 100]) labels_d_fake = tf.placeholder(tf.float32) labels_d_real = tf.placeholder(tf.float32) with tf.variable_scope("gen"): GZ = generator(Z) with tf.variable_scope("disc") as scope: DGZ_raw, DGZ = discriminator(GZ) scope.reuse_variables() DX_raw, DX = discriminator(X) with tf.variable_scope("loss_calculation"): # using Sigmoid Cross Entropy as loss function # discriminator tries to discriminate between X and GZ: # if input to discriminator is X then output prob should be 0 # if input to discriminator is GZ then output prob should be 1 # loss_d_fake = log(1-DGZ_raw) loss_d_fake = tf.nn.sigmoid_cross_entropy_with_logits(labels=labels_d_fake, logits=DGZ_raw) # loss_d_real = log(DX_raw) loss_d_real = tf.nn.sigmoid_cross_entropy_with_logits(labels=labels_d_real, logits=DX_raw) loss_g = tf.nn.sigmoid_cross_entropy_with_logits(labels=tf.ones_like(DGZ_raw), logits=DGZ_raw) # loss_disc = log(1-DGZ_raw) + log(DX_raw) loss_disc = tf.reduce_mean(loss_d_fake + loss_d_real) # gen tries to fool D # that is D should output 0 on seeing GZ # therefore we minimize DGZ-raw(to make it 0) # loss_gen = log(DGZ_raw) loss_gen = tf.reduce_mean(loss_g) GZ_summary = tf.summary.image('GZ', GZ, max_outputs = 10) DGZ_summary = tf.summary.scalar('DGZ', tf.reduce_mean(DGZ)) DX_summary = tf.summary.scalar('DX', tf.reduce_mean(DX)) loss_disc_summary = tf.summary.scalar('loss_disc', loss_disc) loss_gen_summary = tf.summary.scalar('loss_gen', loss_gen) disc_merged = tf.summary.merge([DGZ_summary, DX_summary, loss_disc_summary, GZ_summary]) gen_merged = tf.summary.merge([DGZ_summary, loss_gen_summary]) # Defined Optimizers for both Discriminator and Generator optimizer_disc = tf.train.AdamOptimizer(learning_rate=2e-4, beta1 = 0.5) train_op_disc = optimizer_disc.minimize(loss_disc, var_list=tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, "disc")) optimizer_gen = tf.train.AdamOptimizer(learning_rate=2e-4, beta1 = 0.5) train_op_gen = optimizer_gen.minimize(loss_gen, var_list=tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, "gen")) TAG = 'plotting-graphs_alpha_'+str(flip_alpha)+'_'+DATASET saver = tf.train.Saver(max_to_keep=2) if(DATASET != 'CELEBA'): Xsample_val_set = loader.create_sample_set(100) sess = tf.InteractiveSession() # saver = tf.train.import_meta_graph('saved_analysis_snapshots/snapshots_cifar_complete_from_60_epoch_G_twice/it_2500.ckpt.meta') # saver.restore(sess, 'saved_analysis_snapshots/snapshots_cifar_complete_from_60_epoch_G_twice/it_2500.ckpt') train_writer = tf.summary.FileWriter('logs/' + TAG, sess.graph) if(restore): snapshot_name = tf.train.latest_checkpoint('logs/' + TAG + '/') saver = tf.train.import_meta_graph(snapshot_name + '.meta') print snapshot_name start_iter = int(snapshot_name.split('-')[-1].split('.')[0]) + 1 print start_iter saver.restore(sess, snapshot_name) else: start_iter = 0 sess.run(tf.global_variables_initializer()) num_train_iter = num_train_epochs * (data_size / (batch_size * num_disc_steps)) Z_sample_all = loader.load_batch_Z(100) sum_g_loss, num_g_loss = 0, 0 sum_d_loss, num_d_loss = 0, 0 for iteration in range(start_iter, num_train_iter): for disc_step in range(num_disc_steps): X_feed = loader.load_batch_X(batch_size) Z_feed = loader.load_batch_Z(batch_size) labels_d_real_feed = np.random.choice([0, 1], size=(batch_size,), p=[flip_alpha, 1-flip_alpha]) labels_d_fake_feed = np.ones_like(labels_d_real_feed) - labels_d_real_feed _, cost_disc, DGZ_ret, DX_ret, disc_merged_ret = sess.run(fetches = [train_op_disc, loss_disc, DGZ, DX, disc_merged], feed_dict = {X:X_feed, Z:Z_feed, labels_d_fake: labels_d_fake_feed, labels_d_real:labels_d_real_feed}) sum_d_loss += cost_disc num_d_loss += 1 train_writer.add_summary(disc_merged_ret, iteration) # print ("#%d D #%d\tLossD:%f\tAvgD:%f\tDGZ:%f\tDX:%f" % ((iteration * batch_size * num_disc_steps) / data_size , iteration, cost_disc, sum_d_loss * 1.0 / num_d_loss, np.mean(DGZ_ret), np.mean(DX_ret))) for gen_step in range(num_gen_steps): Z_feed = loader.load_batch_Z(batch_size) _, cost_gen, DGZ_ret, gen_merged_ret = sess.run(fetches = [train_op_gen, loss_gen, DGZ, gen_merged], feed_dict = {Z:Z_feed}) sum_g_loss += cost_gen num_g_loss += 1 train_writer.add_summary(gen_merged_ret, iteration) # print ("#%d G #%d\tLossG:%f\tAvgG:%f\tDGZ:%f" % ((iteration * batch_size * num_disc_steps) / data_size, iteration, cost_gen, sum_g_loss * 1.0 / num_g_loss, np.mean(DGZ_ret))) print ("!%d #%d\tAvgG:%f\tAvgD:%f" % ((iteration * batch_size * num_disc_steps) / data_size, iteration, sum_g_loss * 1.0 / num_g_loss, sum_d_loss * 1.0 / num_d_loss)) if iteration % save_checkpoint_every == 0 and iteration != 0: # save_name = "snapshots/it_%d.ckpt" % iteration if not os.path.isdir('logs/' + TAG): os.makedirs('logs/' + TAG) save_name = 'logs/' + TAG + '/model.ckpt' saver.save(sess, save_name, iteration) print "Snapshot saved to %s" % save_name if iteration % generate_samples_every == 0: if not os.path.isdir('analysis/' + TAG): os.makedirs('analysis/' + TAG) X_sample_all = loader.load_batch_X(100, update_iterator = False) im_samples, im_score = sess.run(fetches=[GZ, DGZ], feed_dict={Z:Z_sample_all}) im_samples = d3_scale(im_samples, out_range=(0,255)) # print "evaluation : %d" % evaluate(Xsample_val_set, im_samples) im_samples = save_sample_images(im_samples, 100) out = sess.run(im_samples) if DATASET != 'CELEBA': out = cv2.cvtColor(out, cv2.COLOR_RGB2BGR) cv2.imwrite("analysis/"+TAG+"/generator_sample_%d.png" % iteration, out) cv2.imwrite("analysis/"+TAG+"/generator_latest.png", out) X_score = sess.run(fetches=[DX], feed_dict={X:X_sample_all}) X_samples = save_sample_images(d3_scale(X_sample_all, out_range=(0,255)), 100) out = sess.run(X_samples) if DATASET != 'CELEBA': out = cv2.cvtColor(out, cv2.COLOR_RGB2BGR) cv2.imwrite("analysis/"+TAG+"/orig_%d.png" % iteration, out) cv2.imwrite("analysis/"+TAG+"/orig_latest.png", out) print "Sample scores: \tDGZ:%f\tDX:%f" % (np.mean(im_score), np.mean(X_score)) train_writer.close() if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument("--dataset", help="dataset name", type=str, choices=['CIFAR10', 'MNIST', 'CELEBA', 'SVHN'], default='MNIST') parser.add_argument("--num_train_epochs", type=int, help="number of training epochs", default=10) parser.add_argument("--num_disc_steps", type=int, help="number of discriminator training steps", default=1) parser.add_argument("--num_gen_steps", type=int, help="number of generator training steps", default=2) parser.add_argument("--batch_size", type=int, help="batch size", default=64) parser.add_argument("--save_checkpoint_every", type=int, help="number of iterations after which a checkpoint is saved", default=250) parser.add_argument("--generate_samples_every", type=int, help="number of iterations after which a sample is generated", default=100) parser.add_argument("--flip_alpha", type=float, help="probability of flipping the label", default=0.3) args = parser.parse_args() main(args.dataset, args.data_size, args.num_train_epochs, args.num_disc_steps, args.num_gen_steps, args.batch_size, args.save_checkpoint_every, args.generate_samples_every, args.flip_alpha)
def main(_): train_images = load_dataset(data_path) fixed_z = np.random.normal(0, 1, [batch_size, 1, 1, z_dim]) z = tf.placeholder(tf.float32, shape=(batch_size, 1, 1, z_dim)) x = tf.placeholder(tf.float32, shape=(batch_size, input_height, input_width, col_dim)) train_bool = tf.placeholder(dtype=tf.bool) #build networks g_x = [generator(z, g_num) for g_num in range(generator_num)] sampler = [ generator(z, g_num, isTraining=False, reuse=True) for g_num in range(generator_num) ] d_real_logits = [discriminator(x, num) for num in range(discrim_num)] d_fake_logits = [ discriminator(sampler[g_num], num, reuse=True) for g_num in range(generator_num) for num in range(discrim_num) ] g_train_logits = [ discriminator(g_x[g_num], num, reuse=True) for g_num in range(generator_num) for num in range(discrim_num) ] #calculate losses d_real_loss = tf.reduce_mean( tf.nn.sigmoid_cross_entropy_with_logits( logits=d_real_logits, labels=tf.ones_like(d_real_logits) * (1.0 - smooth))) d_fake_loss = tf.reduce_mean( tf.nn.sigmoid_cross_entropy_with_logits( logits=d_fake_logits, labels=tf.zeros_like(d_fake_logits))) discrim_loss = d_real_loss + d_fake_loss gen_loss = tf.reduce_mean( tf.nn.sigmoid_cross_entropy_with_logits( logits=g_train_logits, labels=tf.ones_like(g_train_logits))) tf.summary.scalar('Discriminator Loss', discrim_loss) tf.summary.scalar('Generator Loss', gen_loss) #extract trainable variables for g and d var = tf.trainable_variables() d_vars = [[v for v in var if v.name.startswith('discrim_%d' % num)] for num in range(discrim_num)] g_vars = [[v for v in var if v.name.startswith('generator_%d' % num)] for num in range(generator_num)] g_vars.append([v for v in var if ('generator_x') in v.name]) with tf.control_dependencies(tf.get_collection(tf.GraphKeys.UPDATE_OPS)): discrim_optim = tf.train.AdamOptimizer(learning_rate=discrim_lr, beta1=0.5) \ .minimize(discrim_loss, var_list=d_vars) gen_optim = tf.train.AdamOptimizer(learning_rate=gen_lr, beta1=0.5) \ .minimize(gen_loss, var_list=g_vars) gpu_options = tf.GPUOptions(allow_growth=True) config = tf.ConfigProto(gpu_options=gpu_options, log_device_placement=False) merged_summary = tf.summary.merge_all() saver = tf.train.Saver() sess = tf.InteractiveSession() sess.run(tf.global_variables_initializer()) writer = tf.summary.FileWriter(summary_path + str(trial), sess.graph) for v in tf.global_variables(): print(v.name) if (train): np.random.shuffle(train_images) print('Training images:') print(type(train_images)) print(train_images.shape) train_images = train_images.astype('float32') #train_images = train_images/127.5 - 1 #train_images = (train_images - 127.5) / 127.5 train_images = train_images / np.max(train_images) print('Range: ', np.ptp(train_images)) print('Min: ', np.amin(train_images)) print('Max: ', np.amax(train_images)) print('Mean: ', np.mean(train_images)) randGen = np.random.randint(0, generator_num, 9) print(randGen) train_iter = train_size // batch_size train_begin = time.time() d_loss_arr = [None] * train_epoch * train_iter g_loss_arr = [None] * train_epoch * train_iter index = 0 for epoch in range(train_epoch): start_time = time.time() print('Epoch: %d' % epoch) z_gen = np.random.normal(0, 1, [batch_size, 1, 1, z_dim]) for counter in range(train_iter): x_batch = train_images[counter * batch_size:(counter + 1) * batch_size] z_batch = np.random.normal(0, 1, [batch_size, 1, 1, z_dim]) _, d_curr_loss = sess.run([discrim_optim, discrim_loss], feed_dict={ x: x_batch, z: z_batch }) d_loss_arr[index] = d_curr_loss _, g_curr_loss = sess.run([gen_optim, gen_loss], feed_dict={ z: z_batch, x: x_batch }) g_loss_arr[index] = g_curr_loss z_batch = np.random.normal(0, 1, [batch_size, 1, 1, z_dim]) _, g_curr_loss = sess.run([gen_optim, gen_loss], feed_dict={ z: z_batch, x: x_batch }) index += 1 if counter % 200 == 0: print( '\tStep %d: Discriminator loss: %.3f Generator loss: %.3f' % (counter, np.mean(d_loss_arr), np.mean(g_loss_arr))) if counter % 500 == 0: #sample_images(epoch, sess) samples = sess.run(sampler, feed_dict={z: fixed_z}) fig, axes = plt.subplots(3, 3) fig.suptitle('Epoch %d_%d' % (epoch, counter / 500)) for i in range(9): img = np.reshape(samples[randGen[i]][i * 9], (input_height, input_width, col_dim)) axes[i / 3, i % 3].imshow(img) fig.savefig(image_path + 'trial%d/%d_%d.png' % (trial, epoch, counter / 500)) plt.close(fig) saver.save(sess, model_path, global_step=epoch) print('Discriminator loss: %.3f Generator loss: %.3f' % (np.mean(d_loss_arr), np.mean(g_loss_arr))) end_time = time.time() print('Epoch %d took %.2f' % (epoch, end_time - start_time)) train_end = time.time() total_time = train_end - train_begin print('Total training time: %.2f' % total_time) log_trial(total_time) plot_hist(d_loss_arr, g_loss_arr) make_gif()
def train(args): # Context ctx = get_extension_context(args.context, device_id=args.device_id, type_config=args.type_config) nn.set_default_context(ctx) # Args latent = args.latent maps = args.maps batch_size = args.batch_size image_size = args.image_size lambda_ = args.lambda_ # Model # generator loss z = nn.Variable([batch_size, latent]) x_fake = generator(z, maps=maps, up=args.up).apply(persistent=True) p_fake = discriminator(x_fake, maps=maps) loss_gen = gan_loss(p_fake).apply(persistent=True) # discriminator loss p_fake = discriminator(x_fake, maps=maps) x_real = nn.Variable([batch_size, 3, image_size, image_size]) p_real = discriminator(x_real, maps=maps) loss_dis = gan_loss(p_fake, p_real).apply(persistent=True) # gradient penalty eps = F.rand(shape=[batch_size, 1, 1, 1]) x_rmix = eps * x_real + (1.0 - eps) * x_fake p_rmix = discriminator(x_rmix, maps=maps) x_rmix.need_grad = True # Enabling gradient computation for double backward grads = nn.grad([p_rmix], [x_rmix]) l2norms = [F.sum(g**2.0, [1, 2, 3])**0.5 for g in grads] gp = sum([F.mean((l - 1.0)**2.0) for l in l2norms]) loss_dis += lambda_ * gp # generator with fixed value for test z_test = nn.Variable.from_numpy_array(np.random.randn(batch_size, latent)) x_test = generator(z_test, maps=maps, test=True, up=args.up).apply(persistent=True) # Solver solver_gen = S.Adam(args.lrg, args.beta1, args.beta2) solver_dis = S.Adam(args.lrd, args.beta1, args.beta2) with nn.parameter_scope("generator"): params_gen = nn.get_parameters() solver_gen.set_parameters(params_gen) with nn.parameter_scope("discriminator"): params_dis = nn.get_parameters() solver_dis.set_parameters(params_dis) # Monitor monitor = Monitor(args.monitor_path) monitor_loss_gen = MonitorSeries("Generator Loss", monitor, interval=10) monitor_loss_cri = MonitorSeries("Negative Critic Loss", monitor, interval=10) monitor_time = MonitorTimeElapsed("Training Time", monitor, interval=10) monitor_image_tile_train = MonitorImageTile("Image Tile Train", monitor, num_images=batch_size, interval=1, normalize_method=denormalize) monitor_image_tile_test = MonitorImageTile("Image Tile Test", monitor, num_images=batch_size, interval=1, normalize_method=denormalize) # Data Iterator di = data_iterator_cifar10(batch_size, True) # Train loop for i in range(args.max_iter): # Train discriminator x_fake.need_grad = False # no need backward to generator for _ in range(args.n_critic): solver_dis.zero_grad() x_real.d = di.next()[0] / 127.5 - 1.0 z.d = np.random.randn(batch_size, latent) loss_dis.forward(clear_no_need_grad=True) loss_dis.backward(clear_buffer=True) solver_dis.update() # Train generator x_fake.need_grad = True # need backward to generator solver_gen.zero_grad() z.d = np.random.randn(batch_size, latent) loss_gen.forward(clear_no_need_grad=True) loss_gen.backward(clear_buffer=True) solver_gen.update() # Monitor monitor_loss_gen.add(i, loss_gen.d) monitor_loss_cri.add(i, -loss_dis.d) monitor_time.add(i) # Save if i % args.save_interval == 0: monitor_image_tile_train.add(i, x_fake) monitor_image_tile_test.add(i, x_test) nn.save_parameters( os.path.join(args.monitor_path, "params_{}.h5".format(i))) # Last x_test.forward(clear_buffer=True) nn.save_parameters( os.path.join(args.monitor_path, "params_{}.h5".format(i))) monitor_image_tile_train.add(i, x_fake) monitor_image_tile_test.add(i, x_test)
""" graphs """ with tf.device('/gpu:%d' % gpu_id): ''' graph ''' # nodes a_real = tf.placeholder(tf.float32, shape=[None, crop_size, crop_size, 3]) b_real = tf.placeholder(tf.float32, shape=[None, crop_size, crop_size, 3]) a2b_sample = tf.placeholder(tf.float32, shape=[None, crop_size, crop_size, 3]) b2a_sample = tf.placeholder(tf.float32, shape=[None, crop_size, crop_size, 3]) a2b = models.generator(a_real, 'a2b') b2a = models.generator(b_real, 'b2a') b2a2b = models.generator(b2a, 'a2b', reuse=True) a2b2a = models.generator(a2b, 'b2a', reuse=True) a_dis = models.discriminator(a_real, 'a') b2a_dis = models.discriminator(b2a, 'a', reuse=True) b2a_sample_dis = models.discriminator(b2a_sample, 'a', reuse=True) b_dis = models.discriminator(b_real, 'b') a2b_dis = models.discriminator(a2b, 'b', reuse=True) a2b_sample_dis = models.discriminator(a2b_sample, 'b', reuse=True) # losses g_loss_a2b = tf.identity(ops.l2_loss(a2b_dis, tf.ones_like(a2b_dis)), name='g_loss_a2b') g_loss_b2a = tf.identity(ops.l2_loss(b2a_dis, tf.ones_like(b2a_dis)), name='g_loss_b2a') cyc_loss_a = tf.identity(ops.l1_loss(a_real, a2b2a) * 10.0, name='cyc_loss_a') cyc_loss_b = tf.identity(ops.l1_loss(b_real, b2a2b) * 10.0, name='cyc_loss_b') g_loss = g_loss_a2b + g_loss_b2a + cyc_loss_a + cyc_loss_b d_loss_a_real = ops.l2_loss(a_dis, tf.ones_like(a_dis)) d_loss_b2a_sample = ops.l2_loss(b2a_sample_dis, tf.zeros_like(b2a_sample_dis))