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))
