def test_evaluate_g_losses(sess):
_LAMBDA_A = 10
_LAMBDA_B = 10
input_a = tf.random_uniform((5, 7), maxval=1)
cycle_images_a = input_a + 1
input_b = tf.random_uniform((5, 7), maxval=1)
cycle_images_b = input_b - 2
cycle_consistency_loss_a = _LAMBDA_A * losses.cycle_consistency_loss(
real_images=input_a,
generated_images=cycle_images_a,
)
cycle_consistency_loss_b = _LAMBDA_B * losses.cycle_consistency_loss(
real_images=input_b,
generated_images=cycle_images_b,
)
prob_fake_a_is_real = tf.constant([0, 1.0, 0])
prob_fake_b_is_real = tf.constant([1.0, 1.0, 0])
lsgan_loss_a = losses.lsgan_loss_generator(prob_fake_a_is_real)
lsgan_loss_b = losses.lsgan_loss_generator(prob_fake_b_is_real)
assert np.isclose(sess.run(lsgan_loss_a), 0.66666669) and \
np.isclose(sess.run(lsgan_loss_b), 0.3333333) and \
np.isclose(sess.run(cycle_consistency_loss_a), 10) and \
np.isclose(sess.run(cycle_consistency_loss_b), 20)
def compute_losses(self):
"""
In this function we are defining the variables for loss calculations
and training model.
d_loss_A/d_loss_B -> loss for discriminator A/B
g_loss_A/g_loss_B -> loss for generator A/B
*_trainer -> Various trainer for above loss functions
*_summ -> Summary variables for above loss functions
"""
cycle_consistency_loss_a = \
self._lambda_a * losses.cycle_consistency_loss(
real_images=self.input_a, generated_images=self.cycle_images_a,
)
cycle_consistency_loss_b = \
self._lambda_b * losses.cycle_consistency_loss(
real_images=self.input_b, generated_images=self.cycle_images_b,
)
lsgan_loss_a = losses.lsgan_loss_generator(self.prob_fake_a_is_real)
lsgan_loss_b = losses.lsgan_loss_generator(self.prob_fake_b_is_real)
g_loss_A = \
cycle_consistency_loss_a + cycle_consistency_loss_b + lsgan_loss_b
g_loss_B = \
cycle_consistency_loss_b + cycle_consistency_loss_a + lsgan_loss_a
d_loss_A = losses.lsgan_loss_discriminator(
prob_real_is_real=self.prob_real_a_is_real,
prob_fake_is_real=self.prob_fake_pool_a_is_real,
)
d_loss_B = losses.lsgan_loss_discriminator(
prob_real_is_real=self.prob_real_b_is_real,
prob_fake_is_real=self.prob_fake_pool_b_is_real,
)
print g_loss_A, g_loss_B
optimizer = tf.train.AdamOptimizer(self.learning_rate, beta1=0.5)
self.model_vars = tf.trainable_variables()
d_A_vars = [var for var in self.model_vars if 'd_A' in var.name]
g_A_vars = [var for var in self.model_vars if 'g_A' in var.name]
d_B_vars = [var for var in self.model_vars if 'd_B' in var.name]
g_B_vars = [var for var in self.model_vars if 'g_B' in var.name]
self.d_A_trainer = optimizer.minimize(d_loss_A, var_list=d_A_vars)
self.d_B_trainer = optimizer.minimize(d_loss_B, var_list=d_B_vars)
self.g_A_trainer = optimizer.minimize(g_loss_A, var_list=g_A_vars)
self.g_B_trainer = optimizer.minimize(g_loss_B, var_list=g_B_vars)
for var in self.model_vars:
print(var.name)
# Summary variables for tensorboard
self.g_A_loss_summ = tf.summary.scalar("g_A_loss", g_loss_A)
self.g_B_loss_summ = tf.summary.scalar("g_B_loss", g_loss_B)
self.d_A_loss_summ = tf.summary.scalar("d_A_loss", d_loss_A)
self.d_B_loss_summ = tf.summary.scalar("d_B_loss", d_loss_B)
def compute_losses(self):
"""Compute losses."""
self.reconstruction_loss_a = losses.reconstruction_loss(
real_images=self.input_a, generated_images=self.ae_images_a)
self.reconstruction_loss_b = losses.reconstruction_loss(
real_images=self.input_b, generated_images=self.ae_images_b)
self.lsgan_loss_fake_a = losses.lsgan_loss_generator(
self.prob_fake_a_is_real)
self.lsgan_loss_fake_b = losses.lsgan_loss_generator(
self.prob_fake_b_is_real)
self.cycle_consistency_loss_a = losses.cycle_consistency_loss(
real_images=self.input_a, generated_images=self.cycle_images_a)
self.cycle_consistency_loss_b = losses.cycle_consistency_loss(
real_images=self.input_b, generated_images=self.cycle_images_b)
self.g_loss = self._rec_lambda_a * self.reconstruction_loss_a + \
self._rec_lambda_b * self.reconstruction_loss_b + \
self._cycle_lambda_a * self.cycle_consistency_loss_a + \
self._cycle_lambda_b * self.cycle_consistency_loss_b + \
self._lsgan_lambda_a * self.lsgan_loss_fake_a + \
self._lsgan_lambda_b * self.lsgan_loss_fake_b
self.d_loss_A = losses.lsgan_loss_discriminator(
prob_real_is_real=self.prob_real_a_is_real,
prob_fake_is_real=self.prob_fake_pool_a_is_real)
self.d_loss_B = losses.lsgan_loss_discriminator(
prob_real_is_real=self.prob_real_b_is_real,
prob_fake_is_real=self.prob_fake_pool_b_is_real)
self.model_vars = tf.trainable_variables()
d_a_vars = [
var for var in self.model_vars
if 'd1' in var.name or 'd_shared' in var.name
]
d_b_vars = [
var for var in self.model_vars
if 'd2' in var.name or 'd_shared' in var.name
]
g_vars = [
var for var in self.model_vars if 'ae1' in var.name
or 'ae2' in var.name or 'ae_shared' in var.name
]
optimizer = tf.train.AdamOptimizer(self.learning_rate, beta1=0.5)
self.d_A_trainer = optimizer.minimize(self.d_loss_A, var_list=d_a_vars)
self.d_B_trainer = optimizer.minimize(self.d_loss_B, var_list=d_b_vars)
self.g_trainer = optimizer.minimize(self.g_loss, var_list=g_vars)
self.create_summaries()
def test_cycle_consistency_loss_is_none_with_perfect_fakes(sess):
batch_size, height, width, channels = [16, 2, 3, 1]
tf.set_random_seed(0)
images = tf.random_uniform((batch_size, height, width, channels), maxval=1)
loss = losses.cycle_consistency_loss(
real_images=images,
generated_images=images,
)
assert sess.run(loss) == 0
def test_cycle_consistency_loss_is_positive_with_imperfect_fake_x(sess):
batch_size, height, width, channels = [16, 2, 3, 1]
tf.set_random_seed(0)
real_images = tf.random_uniform(
(batch_size, height, width, channels),
maxval=1,
)
generated_images = real_images + 1
loss = losses.cycle_consistency_loss(
real_images=real_images,
generated_images=generated_images,
)
assert sess.run(loss) == 1
def compute_losses(self):
cycle_consistency_loss_a = \
self._lambda_a * losses.cycle_consistency_loss(
real_images=tf.expand_dims(self.input_a[:,:,:,1], axis=3), generated_images=self.cycle_images_a,
)
cycle_consistency_loss_b = \
self._lambda_b * losses.cycle_consistency_loss(
real_images=tf.expand_dims(self.input_b[:,:,:,1], axis=3), generated_images=self.cycle_images_b,
)
lsgan_loss_a = losses.lsgan_loss_generator(self.prob_fake_a_is_real)
lsgan_loss_b = losses.lsgan_loss_generator(self.prob_fake_b_is_real)
lsgan_loss_f = losses.lsgan_loss_generator(self.prob_fea_b_is_real)
lsgan_loss_a_aux = losses.lsgan_loss_generator(
self.prob_fake_a_aux_is_real)
ce_loss_b, dice_loss_b = losses.task_loss(self.pred_mask_fake_b,
self.gt_a)
l2_loss_b = tf.add_n([
0.0001 * tf.nn.l2_loss(v) for v in tf.trainable_variables()
if '/s_B/' in v.name or '/e_B/' in v.name
])
g_loss_A = cycle_consistency_loss_a + cycle_consistency_loss_b + lsgan_loss_b
g_loss_B = cycle_consistency_loss_b + cycle_consistency_loss_a + lsgan_loss_a
self.loss_f_weight = tf.placeholder(tf.float32,
shape=[],
name="loss_f_weight")
self.loss_f_weight_summ = tf.summary.scalar("loss_f_weight",
self.loss_f_weight)
seg_loss_B = ce_loss_b + dice_loss_b + l2_loss_b + 0.1 * g_loss_B + self.loss_f_weight * lsgan_loss_f + 0.1 * lsgan_loss_a_aux
d_loss_A = losses.lsgan_loss_discriminator(
prob_real_is_real=self.prob_real_a_is_real,
prob_fake_is_real=self.prob_fake_pool_a_is_real,
)
d_loss_A_aux = losses.lsgan_loss_discriminator(
prob_real_is_real=self.prob_cycle_a_aux_is_real,
prob_fake_is_real=self.prob_fake_pool_a_aux_is_real,
)
d_loss_A = d_loss_A + d_loss_A_aux
d_loss_B = losses.lsgan_loss_discriminator(
prob_real_is_real=self.prob_real_b_is_real,
prob_fake_is_real=self.prob_fake_pool_b_is_real,
)
d_loss_F = losses.lsgan_loss_discriminator(
prob_real_is_real=self.prob_fea_fake_b_is_real,
prob_fake_is_real=self.prob_fea_b_is_real,
)
optimizer = tf.train.AdamOptimizer(self.learning_rate, beta1=0.5)
optimizer_seg = tf.train.AdamOptimizer(self.learning_rate_seg)
self.model_vars = tf.trainable_variables()
d_A_vars = [var for var in self.model_vars if '/d_A/' in var.name]
d_B_vars = [var for var in self.model_vars if '/d_B/' in var.name]
g_A_vars = [var for var in self.model_vars if '/g_A/' in var.name]
e_B_vars = [var for var in self.model_vars if '/e_B/' in var.name]
de_B_vars = [var for var in self.model_vars if '/de_B/' in var.name]
s_B_vars = [var for var in self.model_vars if '/s_B/' in var.name]
d_F_vars = [var for var in self.model_vars if '/d_F/' in var.name]
self.d_A_trainer = optimizer.minimize(d_loss_A, var_list=d_A_vars)
self.d_B_trainer = optimizer.minimize(d_loss_B, var_list=d_B_vars)
self.g_A_trainer = optimizer.minimize(g_loss_A, var_list=g_A_vars)
self.g_B_trainer = optimizer.minimize(g_loss_B, var_list=de_B_vars)
self.s_B_trainer = optimizer_seg.minimize(seg_loss_B,
var_list=e_B_vars + s_B_vars)
self.d_F_trainer = optimizer.minimize(d_loss_F, var_list=d_F_vars)
for var in self.model_vars:
print(var.name)
# Summary variables for tensorboard
self.g_A_loss_summ = tf.summary.scalar("g_A_loss", g_loss_A)
self.g_B_loss_summ = tf.summary.scalar("g_B_loss", g_loss_B)
self.d_A_loss_summ = tf.summary.scalar("d_A_loss", d_loss_A)
self.d_B_loss_summ = tf.summary.scalar("d_B_loss", d_loss_B)
self.ce_B_loss_summ = tf.summary.scalar("ce_B_loss", ce_loss_b)
self.dice_B_loss_summ = tf.summary.scalar("dice_B_loss", dice_loss_b)
self.l2_B_loss_summ = tf.summary.scalar("l2_B_loss", l2_loss_b)
self.s_B_loss_summ = tf.summary.scalar("s_B_loss", seg_loss_B)
self.s_B_loss_merge_summ = tf.summary.merge([
self.ce_B_loss_summ, self.dice_B_loss_summ, self.l2_B_loss_summ,
self.s_B_loss_summ
])
self.d_F_loss_summ = tf.summary.scalar("d_F_loss", d_loss_F)
def compute_losses(self):
"""
In this function we are defining the variables for loss calculations
and training model.
d_loss_A/d_loss_B -> loss for discriminator A/B
g_loss_A/g_loss_B -> loss for generator A/B
*_trainer -> Various trainer for above loss functions
*_summ -> Summary variables for above loss functions
"""
cycle_consistency_loss_a = \
self._lambda_a * losses.cycle_consistency_loss(
real_images=self.input_a, generated_images=self.cycle_images_a,
)
cycle_consistency_loss_b = \
self._lambda_b * losses.cycle_consistency_loss(
real_images=self.input_b, generated_images=self.cycle_images_b,
)
lsgan_loss_a = losses.lsgan_loss_generator(self.prob_fake_a_is_real)
lsgan_loss_b = losses.lsgan_loss_generator(self.prob_fake_b_is_real)
g_loss_A = \
cycle_consistency_loss_a + cycle_consistency_loss_b + lsgan_loss_b
g_loss_B = \
cycle_consistency_loss_b + cycle_consistency_loss_a + lsgan_loss_a
d_loss_A = losses.lsgan_loss_discriminator(
prob_real_is_real=self.prob_real_a_is_real,
prob_fake_is_real=self.prob_fake_pool_a_is_real,
)
d_loss_B = losses.lsgan_loss_discriminator(
prob_real_is_real=self.prob_real_b_is_real,
prob_fake_is_real=self.prob_fake_pool_b_is_real,
)
optimizer = tf.train.AdamOptimizer(self.learning_rate, beta1=0.5)
optimizer2 = tf.train.AdamOptimizer(self.learning_rate / self.ratio,
beta1=0.5)
self.model_vars = tf.trainable_variables()
d_A_vars = [var for var in self.model_vars if 'd_A' in var.name]
g_A_vars = [var for var in self.model_vars if 'g_A/' in var.name]
d_B_vars = [var for var in self.model_vars if 'd_B' in var.name]
g_B_vars = [var for var in self.model_vars if 'g_B/' in var.name]
g_Ae_vars = [var for var in self.model_vars if 'g_A_ae' in var.name]
g_Be_vars = [var for var in self.model_vars if 'g_B_ae' in var.name]
grads = tf.gradients(g_loss_A, g_A_vars + g_Ae_vars)
grads1 = grads[:len(g_A_vars)]
grads2 = grads[len(g_A_vars):]
train_op1 = optimizer.apply_gradients(zip(grads1, g_A_vars))
train_op2 = optimizer2.apply_gradients(zip(grads2, g_Ae_vars))
self.g_A_trainer = tf.group(train_op1, train_op2)
grads = tf.gradients(g_loss_B, g_B_vars + g_Be_vars)
grads1 = grads[:len(g_B_vars)]
grads2 = grads[len(g_B_vars):]
train_op1_ = optimizer.apply_gradients(zip(grads1, g_B_vars))
train_op2_ = optimizer2.apply_gradients(zip(grads2, g_Be_vars))
self.g_B_trainer = tf.group(train_op1_, train_op2_)
self.d_A_trainer = optimizer.minimize(d_loss_A, var_list=d_A_vars)
self.d_B_trainer = optimizer.minimize(d_loss_B, var_list=d_B_vars)
self.params_ae_c1 = g_A_vars[0]
self.params_ae_c1_B = g_B_vars[0]
for var in self.model_vars:
print(var.name)
# Summary variables for tensorboard
self.g_A_loss_summ = tf.summary.scalar("g_A_loss", g_loss_A)
self.g_B_loss_summ = tf.summary.scalar("g_B_loss", g_loss_B)
self.d_A_loss_summ = tf.summary.scalar("d_A_loss", d_loss_A)
self.d_B_loss_summ = tf.summary.scalar("d_B_loss", d_loss_B)
def compute_losses(self):
cycle_consistency_loss_a = \
self._lambda_a * losses.cycle_consistency_loss(
real_images=self.input_a, generated_images=self.cycle_images_a,
)
cycle_consistency_loss_b = \
self._lambda_b * losses.cycle_consistency_loss(
real_images=self.input_b, generated_images=self.cycle_images_b,
)
lsgan_loss_a = losses.lsgan_loss_generator(self.prob_fake_a_is_real)
lsgan_loss_b = losses.lsgan_loss_generator(self.prob_fake_b_is_real)
lsgan_loss_p = losses.lsgan_loss_generator(
self.prob_pred_mask_b_is_real)
lsgan_loss_p_ll = losses.lsgan_loss_generator(
self.prob_pred_mask_b_ll_is_real)
lsgan_loss_a_aux = losses.lsgan_loss_generator(
self.prob_fake_a_aux_is_real)
ce_loss_b, dice_loss_b = losses.task_loss(self.pred_mask_fake_b,
self.gt_a)
ce_loss_b_ll, dice_loss_b_ll = losses.task_loss(
self.pred_mask_fake_b_ll, self.gt_a)
l2_loss_b = tf.add_n([
0.0001 * tf.nn.l2_loss(v) for v in tf.trainable_variables()
if '/s_B/' in v.name or '/s_B_ll/' in v.name or '/e_B/' in v.name
])
g_loss_A = cycle_consistency_loss_a + cycle_consistency_loss_b + lsgan_loss_b
g_loss_B = cycle_consistency_loss_b + cycle_consistency_loss_a + lsgan_loss_a
seg_loss_B = ce_loss_b + dice_loss_b + l2_loss_b + 0.1 * (
ce_loss_b_ll + dice_loss_b_ll
) + 0.1 * g_loss_B + 0.1 * lsgan_loss_p + 0.01 * lsgan_loss_p_ll + 0.1 * lsgan_loss_a_aux
d_loss_A = losses.lsgan_loss_discriminator(
prob_real_is_real=self.prob_real_a_is_real,
prob_fake_is_real=self.prob_fake_pool_a_is_real,
)
d_loss_A_aux = losses.lsgan_loss_discriminator(
prob_real_is_real=self.prob_cycle_a_aux_is_real,
prob_fake_is_real=self.prob_fake_pool_a_aux_is_real,
)
d_loss_A = d_loss_A + d_loss_A_aux
d_loss_B = losses.lsgan_loss_discriminator(
prob_real_is_real=self.prob_real_b_is_real,
prob_fake_is_real=self.prob_fake_pool_b_is_real,
)
d_loss_P = losses.lsgan_loss_discriminator(
prob_real_is_real=self.prob_pred_mask_fake_b_is_real,
prob_fake_is_real=self.prob_pred_mask_b_is_real,
)
d_loss_P_ll = losses.lsgan_loss_discriminator(
prob_real_is_real=self.prob_pred_mask_fake_b_ll_is_real,
prob_fake_is_real=self.prob_pred_mask_b_ll_is_real,
)
optimizer_gan = tf.train.AdamOptimizer(self.learning_rate_gan,
beta1=0.5)
optimizer_seg = tf.train.AdamOptimizer(self.learning_rate_seg)
self.model_vars = tf.trainable_variables()
d_A_vars = [var for var in self.model_vars if '/d_A/' in var.name]
d_B_vars = [var for var in self.model_vars if '/d_B/' in var.name]
g_A_vars = [var for var in self.model_vars if '/g_A/' in var.name]
e_B_vars = [var for var in self.model_vars if '/e_B/' in var.name]
de_B_vars = [var for var in self.model_vars if '/de_B/' in var.name]
s_B_vars = [var for var in self.model_vars if '/s_B/' in var.name]
s_B_ll_vars = [
var for var in self.model_vars if '/s_B_ll/' in var.name
]
d_P_vars = [var for var in self.model_vars if '/d_P/' in var.name]
d_P_ll_vars = [
var for var in self.model_vars if '/d_P_ll/' in var.name
]
self.d_A_trainer = optimizer_gan.minimize(d_loss_A, var_list=d_A_vars)
self.d_B_trainer = optimizer_gan.minimize(d_loss_B, var_list=d_B_vars)
self.g_A_trainer = optimizer_gan.minimize(g_loss_A, var_list=g_A_vars)
self.g_B_trainer = optimizer_gan.minimize(g_loss_B, var_list=de_B_vars)
self.d_P_trainer = optimizer_gan.minimize(d_loss_P, var_list=d_P_vars)
self.d_P_ll_trainer = optimizer_gan.minimize(d_loss_P_ll,
var_list=d_P_ll_vars)
self.s_B_trainer = optimizer_seg.minimize(seg_loss_B,
var_list=e_B_vars +
s_B_vars + s_B_ll_vars)
for var in self.model_vars:
print(var.name)
# Summary variables for tensorboard
self.g_A_loss_summ = tf.summary.scalar("g_A_loss", g_loss_A)
self.g_B_loss_summ = tf.summary.scalar("g_B_loss", g_loss_B)
self.d_A_loss_summ = tf.summary.scalar("d_A_loss", d_loss_A)
self.d_B_loss_summ = tf.summary.scalar("d_B_loss", d_loss_B)
self.ce_B_loss_summ = tf.summary.scalar("ce_B_loss", ce_loss_b)
self.dice_B_loss_summ = tf.summary.scalar("dice_B_loss", dice_loss_b)
self.l2_B_loss_summ = tf.summary.scalar("l2_B_loss", l2_loss_b)
self.s_B_loss_summ = tf.summary.scalar("s_B_loss", seg_loss_B)
self.s_B_loss_merge_summ = tf.summary.merge([
self.ce_B_loss_summ, self.dice_B_loss_summ, self.l2_B_loss_summ,
self.s_B_loss_summ
])
self.d_P_loss_summ = tf.summary.scalar("d_P_loss", d_loss_P)
self.d_P_ll_loss_summ = tf.summary.scalar("d_P_loss_ll", d_loss_P_ll)
self.d_P_loss_merge_summ = tf.summary.merge(
[self.d_P_loss_summ, self.d_P_ll_loss_summ])
def compute_losses(self):
"""
In this function we are defining the variables for loss calculations
and training model.
d_loss_A/d_loss_B -> loss for discriminator A/B
g_loss_A/g_loss_B -> loss for generator A/B
seg_cost -> Segmentation loss in target domain
*_trainer -> Various trainer for above loss functions
*_summ -> Summary variables for above loss functions
"""
self.seg_cost = tf.reduce_mean(tf.square(self.manual_seg
- self.prediction))
cycle_consistency_loss_a = self._lambda_a \
* losses.cycle_consistency_loss(real_images=self.input_a,
generated_images=self.cycle_images_a)
cycle_consistency_loss_b = self._lambda_b \
* losses.cycle_consistency_loss(real_images=self.input_b,
generated_images=self.cycle_images_b)
lsgan_loss_a = \
losses.lsgan_loss_generator(self.prob_fake_a_is_real)
lsgan_loss_b = \
losses.lsgan_loss_generator(self.prob_fake_b_is_real)
g_loss_A = cycle_consistency_loss_a + cycle_consistency_loss_b \
+ lsgan_loss_b + 10 * self.seg_cost
g_loss_B = cycle_consistency_loss_b + cycle_consistency_loss_a \
+ lsgan_loss_a
d_loss_A = \
losses.lsgan_loss_discriminator(prob_real_is_real=self.prob_real_a_is_real,
prob_fake_is_real=self.prob_fake_pool_a_is_real)
d_loss_B = \
losses.lsgan_loss_discriminator(prob_real_is_real=self.prob_real_b_is_real,
prob_fake_is_real=self.prob_fake_pool_b_is_real)
optimizer = tf.train.AdamOptimizer(self.learning_rate,
beta1=0.5)
self.model_vars = tf.trainable_variables()
d_A_vars = [var for var in self.model_vars if 'd_A' in var.name]
g_A_vars = [var for var in self.model_vars if 'g_A/'
in var.name]
d_B_vars = [var for var in self.model_vars if 'd_B' in var.name]
g_B_vars = [var for var in self.model_vars if 'g_B/'
in var.name]
g_Ae_vars = [var for var in self.model_vars if 'g_A_ae'
in var.name]
g_Be_vars = [var for var in self.model_vars if 'g_B_ae'
in var.name]
seg_vars = [var for var in self.model_vars if 'seg_net'
in var.name]
# self.g_A_trainer = optimizer.minimize(g_loss_A, var_list=g_A_vars+g_Ae_vars)
self.g_A_trainer = optimizer.minimize(g_loss_A,
var_list=g_A_vars + g_Ae_vars + seg_vars)
self.g_B_trainer = optimizer.minimize(g_loss_B,
var_list=g_B_vars + g_Be_vars)
self.g_A_trainer_bis = optimizer.minimize(g_loss_A,
var_list=g_A_vars + seg_vars)
self.g_B_trainer_bis = optimizer.minimize(g_loss_B,
var_list=g_B_vars)
self.d_A_trainer = optimizer.minimize(d_loss_A,
var_list=d_A_vars)
self.d_B_trainer = optimizer.minimize(d_loss_B,
var_list=d_B_vars)
self.seg_trainer = optimizer.minimize(self.seg_cost,
var_list=seg_vars)
self.params_ae_c1 = g_A_vars[0]
self.params_ae_c1_B = g_B_vars[0]
for var in self.model_vars:
print var.name
# Summary variables for tensorboard
self.g_A_loss_summ = tf.summary.scalar('g_A_loss', g_loss_A)
self.g_B_loss_summ = tf.summary.scalar('g_B_loss', g_loss_B)
self.d_A_loss_summ = tf.summary.scalar('d_A_loss', d_loss_A)
self.d_B_loss_summ = tf.summary.scalar('d_B_loss', d_loss_B)
self.segmentation_loss_sum = \
tf.summary.scalar('segmentation_loss', self.seg_cost)
def compute_losses(self):
"""
In this function we are defining the variables for loss calculations
and training model.
d_loss_A/d_loss_B -> loss for discriminator A/B
g_loss_A/g_loss_B -> loss for generator A/B
*_trainer -> Various trainer for above loss functions
*_summ -> Summary variables for above loss functions
"""
cycle_consistency_loss_a = self._lambda_a * losses.cycle_consistency_loss(real_images=self.input_a, generated_images=self.cycle_images_a)
cycle_consistency_loss_b = self._lambda_b * losses.cycle_consistency_loss(real_images=self.input_b, generated_images=self.cycle_images_b)
ssim_loss_A = 0.25 * (2 - (tf.reduce_mean(tf.image.ssim_multiscale(self.input_a[:,:,:,:-1], self.fake_images_b,2,power_factors=[0.0448, 0.2856, 0.3001]) + \
tf.image.ssim_multiscale(self.fake_images_b, self.cycle_images_a[:,:,:,:-1],2,power_factors=[0.0448, 0.2856, 0.3001]))))
ssim_loss_B = 0.25 * (2 - (tf.reduce_mean(tf.image.ssim_multiscale(self.input_b, self.fake_images_a[:,:,:,:-1],2,power_factors=[0.0448, 0.2856, 0.3001]) + \
tf.image.ssim_multiscale(self.fake_images_a[:,:,:,:-1], self.cycle_images_b,2,power_factors=[0.0448, 0.2856, 0.3001]))))
lsgan_loss_a = losses.lsgan_loss_generator(self.prob_fake_a_is_real)
lsgan_loss_b = losses.lsgan_loss_generator(self.prob_fake_b_is_real)
grad_loss_B = tf.reduce_mean(tf.image.image_gradients(tf.expand_dims(self.fake_images_a[:,:,:,-1], axis = 3)))
g_loss_A = cycle_consistency_loss_a + cycle_consistency_loss_b + lsgan_loss_b + ssim_loss_A + ssim_loss_B
g_loss_B = cycle_consistency_loss_b + cycle_consistency_loss_a + lsgan_loss_a + ssim_loss_B + ssim_loss_A + grad_loss_B
#g_loss_A = cycle_consistency_loss_a + cycle_consistency_loss_b + lsgan_loss_b
#g_loss_B = cycle_consistency_loss_b + cycle_consistency_loss_a + lsgan_loss_a
d_loss_A = losses.lsgan_loss_discriminator(
prob_real_is_real=self.prob_real_a_is_real,
prob_fake_is_real=self.prob_fake_pool_a_is_real,
)
d_loss_B = losses.lsgan_loss_discriminator(
prob_real_is_real=self.prob_real_b_is_real,
prob_fake_is_real=self.prob_fake_pool_b_is_real,
)
optimizer = tf.train.AdamOptimizer(self.learning_rate, beta1=0.9)
self.model_vars = tf.trainable_variables()
d_A_vars = [var for var in self.model_vars if 'd_A' in var.name]
g_A_vars = [var for var in self.model_vars if 'g_A' in var.name]
d_B_vars = [var for var in self.model_vars if 'd_B' in var.name]
g_B_vars = [var for var in self.model_vars if 'g_B' in var.name]
self.d_A_trainer = optimizer.minimize(d_loss_A, var_list=d_A_vars)
self.d_B_trainer = optimizer.minimize(d_loss_B, var_list=d_B_vars)
self.g_A_trainer = optimizer.minimize(g_loss_A, var_list=g_A_vars)
self.g_B_trainer = optimizer.minimize(g_loss_B, var_list=g_B_vars)
for var in self.model_vars:
print(var.name)
# Summary variables for tensorboard
self.g_A_loss_summ = tf.summary.scalar("g_A_loss", g_loss_A)
self.g_B_loss_summ = tf.summary.scalar("g_B_loss", g_loss_B)
self.d_A_loss_summ = tf.summary.scalar("d_A_loss", d_loss_A)
self.d_B_loss_summ = tf.summary.scalar("d_B_loss", d_loss_B)
self.ssim_A_loss_summ = tf.summary.scalar("ssim_A_loss", ssim_loss_A)
self.ssim_B_loss_summ = tf.summary.scalar("ssim_B_loss", ssim_loss_B)
self.grad_B_loss_summ = tf.summary.scalar("gradient_loss", grad_loss_B)
def _buildGraph(self):
# placeholder variables for feeding
real_a = tf.placeholder(
tf.float32,
shape=[
None, # enables variable batch size
self.feature_size
],
name="A")
real_b = tf.placeholder(
tf.float32,
shape=[
None, # same feature size for both domains
self.feature_size
],
name="B")
dropout = tf.placeholder_with_default(1., shape=[], name="dropout")
train_status = tf.placeholder_with_default(True,
shape=[],
name="train_status")
# random intial values for fake pools
fake_pool_a_sample = tf.placeholder(
tf.float32,
shape=[
None, # enables variable batch size
self.feature_size
],
name="fake_pool_a_sample")
fake_pool_b_sample = tf.placeholder(
tf.float32,
shape=[
None, # same feature size for both domains
self.feature_size
],
name="fake_pool_b_sample")
# construct generator and discriminator functions as class instances
G_a2b = layers.generator(name='G_a2b',
architecture=self.architecture['G_a2b'],
activation=tf.nn.leaky_relu,
output_activation=tf.nn.leaky_relu,
dropout=dropout)
# linear discriminator
D_a = layers.discriminator(name='D_a',
architecture=self.architecture['D_a'],
activation=None,
dropout=dropout)
G_b2a = layers.generator(name='G_b2a',
architecture=self.architecture['G_b2a'],
activation=tf.nn.leaky_relu,
output_activation=tf.nn.leaky_relu,
dropout=dropout)
# linear discriminator
D_b = layers.discriminator(name='D_b',
architecture=self.architecture['D_b'],
activation=None,
dropout=dropout)
# print generator and discriminator architecture
for f in [G_a2b, D_a, G_b2a, D_b]:
f.print_architecture()
# generate fake and cycle profiles
fake_a = G_b2a(real_b, is_train=train_status)
fake_b = G_a2b(real_a, is_train=train_status)
cycle_a = G_b2a(fake_b, is_train=train_status)
cycle_b = G_a2b(fake_a, is_train=train_status)
# construct recent fake profile pool
"""
See Zhu et al., "Training Details":
to reduce model oscillation [15],
we follow Shrivastava et al.’s strategy [46] and update the
discriminators using a history of generated images rather than
the ones produced by the latest generators. We keep an image
buffer that stores the 50 previously created images.
"""
# get discriminator predictions
prob_real_a_is_a = D_a(real_a, is_train=train_status)
prob_fake_a_is_a = D_a(fake_a, is_train=train_status)
prob_real_b_is_b = D_b(real_b, is_train=train_status)
prob_fake_b_is_b = D_b(fake_b, is_train=train_status)
prob_fake_pool_a_is_a = D_a(fake_pool_a_sample, is_train=train_status)
prob_fake_pool_b_is_b = D_b(fake_pool_b_sample, is_train=train_status)
# calculate losses
## Generator loss
lsgan_loss_a = losses.lsgan_loss_generator(prob_fake_a_is_a)
lsgan_loss_b = losses.lsgan_loss_generator(prob_fake_b_is_b)
## Cycle-consistency loss
## Weighted by quality of the generated image
cycle_consistency_loss_a =\
self.lambda_a * losses.cycle_consistency_loss(
real_images=real_a, generated_images=cycle_a) *\
tf.maximum(0.,tf.reduce_mean(prob_fake_b_is_b))
cycle_consistency_loss_b =\
self.lambda_b * losses.cycle_consistency_loss(
real_images=real_b, generated_images=cycle_b) *\
tf.maximum(0.,tf.reduce_mean(prob_fake_a_is_a))
g_loss_a2b =\
cycle_consistency_loss_a + cycle_consistency_loss_b + lsgan_loss_b*2
g_loss_b2a =\
cycle_consistency_loss_b + cycle_consistency_loss_a + lsgan_loss_a*2
## Discriminator loss
d_loss_a = losses.lsgan_loss_discriminator(
prob_real_is_real=prob_real_a_is_a,
prob_fake_is_real=prob_fake_pool_a_is_a,
)
d_loss_b = losses.lsgan_loss_discriminator(
prob_real_is_real=prob_real_b_is_b,
prob_fake_is_real=prob_fake_pool_b_is_b,
)
# construct optimizers
global_step = tf.Variable(0, trainable=False)
# make learning rate adjustable
curr_lr_ = tf.Variable(self.learning_rate, trainable=False)
optimizer = tf.train.AdamOptimizer(curr_lr_, beta1=0.5)
model_vars = tf.trainable_variables()
g_a2b_vars = [var for var in model_vars if 'G_a2b' in var.name]
d_b_vars = [var for var in model_vars if 'D_b' in var.name]
g_b2a_vars = [var for var in model_vars if 'G_b2a' in var.name]
d_a_vars = [var for var in model_vars if 'D_a' in var.name]
g_a2b_trainer = optimizer.minimize(g_loss_a2b, var_list=g_a2b_vars)
d_b_trainer = optimizer.minimize(d_loss_b, var_list=d_b_vars)
g_b2a_trainer = optimizer.minimize(g_loss_b2a, var_list=g_b2a_vars)
d_a_trainer = optimizer.minimize(d_loss_a,
var_list=d_a_vars,
global_step=global_step)
for var in model_vars:
print(var.name)
# Summary variables for tensorboard
tf.summary.scalar("D_A(x)", tf.reduce_mean(prob_fake_a_is_a))
tf.summary.scalar("D_B(x)", tf.reduce_mean(prob_fake_b_is_b))
tf.summary.scalar("G_A2B_loss", g_loss_a2b)
tf.summary.scalar("G_B2A_loss", g_loss_b2a)
tf.summary.scalar("D_A_loss", d_loss_a)
tf.summary.scalar("D_B_loss", d_loss_b)
# Merge summary
merged_summary = tf.summary.merge_all()
# return handles
return (real_a, real_b, dropout, train_status, fake_a, fake_b, cycle_a,
cycle_b, fake_pool_a_sample, fake_pool_b_sample, g_loss_a2b,
d_loss_b, g_loss_b2a, d_loss_a, g_a2b_trainer, d_b_trainer,
g_b2a_trainer, d_a_trainer, curr_lr_, global_step,
merged_summary)
