def build_graph(args, a_r, b_r, a2b_s, b2a_s):
with tf.device('/gpu:{}'.format(args.gpu)):
a2b = g_net(a_r, 'a2b')
b2a = g_net(b_r, 'b2a')
a2b2a = g_net(a2b, 'b2a', reuse=True)
b2a2b = g_net(b2a, 'a2b', reuse=True)
cvt = (a2b, b2a, a2b2a, b2a2b)
a_d = d_net(a_r, 'a')
b2a_d = d_net(b2a, 'a', reuse=True)
b2a_s_d = d_net(b2a_s, 'a', reuse=True)
b_d = d_net(b_r, 'b')
a2b_d = d_net(a2b, 'b', reuse=True)
a2b_s_d = d_net(a2b_s, 'b', reuse=True)
g_loss_a2b = tf.identity(ops.l2_loss(a2b_d, tf.ones_like(a2b_d)),
name='g_loss_a2b')
g_loss_b2a = tf.identity(ops.l2_loss(b2a_d, tf.ones_like(b2a_d)),
name='g_loss_b2a')
cyc_loss_a = tf.identity(ops.l1_loss(a_r, a2b2a) * 10.0,
name='cyc_loss_a')
cyc_loss_b = tf.identity(ops.l1_loss(b_r, b2a2b) * 10.0,
name='cyc_loss_b')
g_loss = g_loss_a2b + g_loss_b2a + cyc_loss_a + cyc_loss_b
d_loss_a_r = ops.l2_loss(a_d, tf.ones_like(a_d))
d_loss_b2a_s = ops.l2_loss(b2a_s_d, tf.zeros_like(b2a_s_d))
d_loss_a = tf.identity((d_loss_a_r + d_loss_b2a_s) / 2.,
name='d_loss_a')
d_loss_b_r = ops.l2_loss(b_d, tf.ones_like(b_d))
d_loss_a2b_s = ops.l2_loss(a2b_s_d, tf.zeros_like(a2b_s_d))
d_loss_b = tf.identity((d_loss_b_r + d_loss_a2b_s) / 2.,
name='d_loss_b')
g_sum = ops.summary_tensors(
[g_loss_a2b, g_loss_b2a, cyc_loss_a, cyc_loss_b])
d_sum_a = ops.summary(d_loss_a)
d_sum_b = ops.summary(d_loss_b)
sum_ = (g_sum, d_sum_a, d_sum_b)
all_var = tf.trainable_variables()
g_var = [
var for var in all_var
if 'a2b_g' in var.name or 'b2a_g' in var.name
]
d_a_var = [var for var in all_var if 'a_d' in var.name]
d_b_var = [var for var in all_var if 'b_d' in var.name]
g_tr_op = tf.train.AdamOptimizer(args.lr, beta1=args.beta1).minimize(
g_loss, var_list=g_var)
d_tr_op_a = tf.train.AdamOptimizer(args.lr, beta1=args.beta1).minimize(
d_loss_a, var_list=d_a_var)
d_tr_op_b = tf.train.AdamOptimizer(args.lr, beta1=args.beta1).minimize(
d_loss_b, var_list=d_b_var)
tr_op = (g_tr_op, d_tr_op_a, d_tr_op_b)
return cvt, sum_, tr_op
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))
d_loss_a = tf.identity((d_loss_a_real + d_loss_b2a_sample) / 2.0,
name='d_loss_a')
d_loss_b_real = ops.l2_loss(b_dis, tf.ones_like(b_dis))
d_loss_a2b_sample = ops.l2_loss(a2b_sample_dis,
tf.zeros_like(a2b_sample_dis))
d_loss_b = tf.identity((d_loss_b_real + d_loss_a2b_sample) / 2.0,
name='d_loss_b')
# summaries
g_summary = ops.summary_tensors(
[g_loss_a2b, g_loss_b2a, cyc_loss_a, cyc_loss_b])
d_summary_a = ops.summary(d_loss_a)
d_summary_b = ops.summary(d_loss_b)
''' optim '''
t_var = tf.trainable_variables()
d_a_var = [var for var in t_var if 'a_discriminator' in var.name]
d_b_var = [var for var in t_var if 'b_discriminator' in var.name]
g_var = [
var for var in t_var
if 'a2b_generator' in var.name or 'b2a_generator' in var.name
]
d_a_train_op = tf.train.AdamOptimizer(lr,
beta1=0.5).minimize(d_loss_a,
var_list=d_a_var)
d_b_train_op = tf.train.AdamOptimizer(lr,
beta1=0.5).minimize(d_loss_b,
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))
d_loss_a = tf.identity((d_loss_a_real + d_loss_b2a_sample) / 2.0, name='d_loss_a')
d_loss_b_real = ops.l2_loss(b_dis, tf.ones_like(b_dis))
d_loss_a2b_sample = ops.l2_loss(a2b_sample_dis, tf.zeros_like(a2b_sample_dis))
d_loss_b = tf.identity((d_loss_b_real + d_loss_a2b_sample) / 2.0, name='d_loss_b')
# summaries
g_summary = ops.summary_tensors([g_loss_a2b, g_loss_b2a, cyc_loss_a, cyc_loss_b])
d_summary_a = ops.summary(d_loss_a)
d_summary_b = ops.summary(d_loss_b)
''' optim '''
t_var = tf.trainable_variables()
d_a_var = [var for var in t_var if 'a_discriminator' in var.name]
d_b_var = [var for var in t_var if 'b_discriminator' in var.name]
g_var = [var for var in t_var if 'a2b_generator' in var.name or 'b2a_generator' in var.name]
d_a_train_op = tf.train.AdamOptimizer(lr, beta1=0.5).minimize(d_loss_a, var_list=d_a_var)
d_b_train_op = tf.train.AdamOptimizer(lr, beta1=0.5).minimize(d_loss_b, var_list=d_b_var)
g_train_op = tf.train.AdamOptimizer(lr, beta1=0.5).minimize(g_loss, var_list=g_var)
""" train """
''' init '''