def build_model(self):
# Placeholdersn
if self.config.trainer.init_type == "normal": ## test
self.init_kernel = tf.random_normal_initializer(mean=0.0,
stddev=0.02)
elif self.config.trainer.init_type == "xavier":
self.init_kernel = tf.contrib.layers.xavier_initializer(
uniform=False, seed=None, dtype=tf.float32)
self.is_training = tf.placeholder(tf.bool)
self.image_tensor = tf.placeholder(tf.float32,
shape=[None] +
self.config.trainer.image_dims,
name="x")
self.noise_tensor = tf.placeholder(
tf.float32,
shape=[None, self.config.trainer.noise_dim],
name="noise")
self.true_labels = tf.placeholder(dtype=tf.float32,
shape=[None, 1],
name="true_labels")
self.generated_labels = tf.placeholder(dtype=tf.float32,
shape=[None, 1],
name="gen_labels")
self.real_noise = tf.placeholder(dtype=tf.float32,
shape=[None] +
self.config.trainer.image_dims,
name="real_noise")
self.fake_noise = tf.placeholder(dtype=tf.float32,
shape=[None] +
self.config.trainer.image_dims,
name="fake_noise")
# Building the Graph
with tf.variable_scope("ALAD"):
# Generated noise from the encoder
with tf.variable_scope("Encoder_Model"):
self.z_gen = self.encoder(
self.image_tensor,
do_spectral_norm=self.config.trainer.do_spectral_norm)
# Generated image and reconstructed image from the Generator
with tf.variable_scope("Generator_Model"):
self.img_gen = self.generator(
self.noise_tensor) + self.fake_noise
self.rec_img = self.generator(self.z_gen)
# Reconstructed image of generated image from the encoder
with tf.variable_scope("Encoder_Model"):
self.rec_z = self.encoder(
self.img_gen, do_spectral_norm=self.config.spectral_norm)
# Discriminator results of (G(z),z) and (x, E(x))
with tf.variable_scope("Discriminator_Model_XZ"):
l_generator, inter_layer_rct_xz = self.discriminator_xz(
self.img_gen,
self.noise_tensor,
do_spectral_norm=self.config.spectral_norm)
l_encoder, inter_layer_inp_xz = self.discriminator_xz(
self.image_tensor + self.real_noise,
self.z_gen,
do_spectral_norm=self.config.do_spectral_norm,
)
# Discrimeinator results of (x, x) and (x, G(E(x))
with tf.variable_scope("Discriminator_Model_XX"):
x_logit_real, inter_layer_inp_xx = self.discriminator_xx(
self.image_tensor + self.real_noise,
self.image_tensor + self.real_noise,
do_spectral_norm=self.config.spectral_norm,
)
x_logit_fake, inter_layer_rct_xx = self.discriminator_xx(
self.image_tensor + self.real_noise,
self.rec_img,
do_spectral_norm=self.config.spectral_norm,
)
# Discriminator results of (z, z) and (z, E(G(z))
with tf.variable_scope("Discriminator_Model_ZZ"):
z_logit_real, _ = self.discriminator_zz(
self.noise_tensor,
self.noise_tensor,
do_spectral_norm=self.config.spectral_norm)
z_logit_fake, _ = self.discriminator_zz(
self.noise_tensor,
self.rec_z,
do_spectral_norm=self.config.spectral_norm)
########################################################################
# LOSS FUNCTIONS
########################################################################
with tf.name_scope("Loss_Functions"):
# discriminator xz
# Discriminator should classify encoder pair as real
loss_dis_enc = tf.reduce_mean(
tf.nn.sigmoid_cross_entropy_with_logits(
labels=self.true_labels, logits=l_encoder))
# Discriminator should classify generator pair as fake
loss_dis_gen = tf.reduce_mean(
tf.nn.sigmoid_cross_entropy_with_logits(
labels=self.generated_labels, logits=l_generator))
self.dis_loss_xz = loss_dis_gen + loss_dis_enc
# discriminator xx
x_real_dis = tf.nn.sigmoid_cross_entropy_with_logits(
logits=x_logit_real, labels=tf.ones_like(x_logit_real))
x_fake_dis = tf.nn.sigmoid_cross_entropy_with_logits(
logits=x_logit_fake, labels=tf.zeros_like(x_logit_fake))
self.dis_loss_xx = tf.reduce_mean(x_real_dis + x_fake_dis)
# discriminator zz
z_real_dis = tf.nn.sigmoid_cross_entropy_with_logits(
logits=z_logit_real, labels=tf.ones_like(z_logit_real))
z_fake_dis = tf.nn.sigmoid_cross_entropy_with_logits(
logits=z_logit_fake, labels=tf.zeros_like(z_logit_fake))
self.dis_loss_zz = tf.reduce_mean(z_real_dis + z_fake_dis)
# Compute the whole discriminator loss
self.loss_discriminator = (self.dis_loss_xz + self.dis_loss_xx +
self.dis_loss_zz
if self.config.trainer.allow_zz else
self.dis_loss_xz + self.dis_loss_xx)
# generator and encoder
if self.config.trainer.flip_labels:
labels_gen = tf.zeros_like(l_generator)
labels_enc = tf.ones_like(l_encoder)
else:
labels_gen = tf.ones_like(l_generator)
labels_enc = tf.zeros_like(l_encoder)
gen_loss_xz = tf.reduce_mean(
tf.nn.sigmoid_cross_entropy_with_logits(labels=labels_gen,
logits=l_generator))
enc_loss_xz = tf.reduce_mean(
tf.nn.sigmoid_cross_entropy_with_logits(labels=labels_enc,
logits=l_encoder))
x_real_gen = tf.nn.sigmoid_cross_entropy_with_logits(
logits=x_logit_real, labels=tf.zeros_like(x_logit_real))
x_fake_gen = tf.nn.sigmoid_cross_entropy_with_logits(
logits=x_logit_fake, labels=tf.ones_like(x_logit_fake))
z_real_gen = tf.nn.sigmoid_cross_entropy_with_logits(
logits=z_logit_real, labels=tf.zeros_like(z_logit_real))
z_fake_gen = tf.nn.sigmoid_cross_entropy_with_logits(
logits=z_logit_fake, labels=tf.ones_like(z_logit_fake))
cost_x = tf.reduce_mean(x_real_gen + x_fake_gen)
cost_z = tf.reduce_mean(z_real_gen + z_fake_gen)
cycle_consistency_loss = cost_x + cost_z if self.config.trainer.allow_zz else cost_x
self.loss_generator = gen_loss_xz + cycle_consistency_loss
self.loss_encoder = enc_loss_xz + cycle_consistency_loss
########################################################################
# OPTIMIZATION
########################################################################
with tf.name_scope("Optimizers"):
# control op dependencies for batch norm and trainable variables
all_variables = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES)
self.dxzvars = [
v for v in all_variables
if v.name.startswith("ALAD/Discriminator_Model_XZ")
]
self.dxxvars = [
v for v in all_variables
if v.name.startswith("ALAD/Discriminator_Model_XX")
]
self.dzzvars = [
v for v in all_variables
if v.name.startswith("ALAD/Discriminator_Model_ZZ")
]
self.gvars = [
v for v in all_variables
if v.name.startswith("ALAD/Generator_Model")
]
self.evars = [
v for v in all_variables
if v.name.startswith("ALAD/Encoder_Model")
]
self.update_ops_gen = tf.get_collection(
tf.GraphKeys.UPDATE_OPS, scope="ALAD/Generator_Model")
self.update_ops_enc = tf.get_collection(tf.GraphKeys.UPDATE_OPS,
scope="ALAD/Encoder_Model")
self.update_ops_dis_xz = tf.get_collection(
tf.GraphKeys.UPDATE_OPS, scope="ALAD/Discriminator_Model_XZ")
self.update_ops_dis_xx = tf.get_collection(
tf.GraphKeys.UPDATE_OPS, scope="ALAD/Discriminator_Model_XX")
self.update_ops_dis_zz = tf.get_collection(
tf.GraphKeys.UPDATE_OPS, scope="ALAD/Discriminator_Model_ZZ")
self.disc_optimizer = tf.train.AdamOptimizer(
learning_rate=self.config.trainer.discriminator_l_rate,
beta1=self.config.trainer.optimizer_adam_beta1,
beta2=self.config.trainer.optimizer_adam_beta2,
)
self.gen_optimizer = tf.train.AdamOptimizer(
learning_rate=self.config.trainer.generator_l_rate,
beta1=self.config.trainer.optimizer_adam_beta1,
beta2=self.config.trainer.optimizer_adam_beta2,
)
self.enc_optimizer = tf.train.AdamOptimizer(
learning_rate=self.config.trainer.generator_l_rate,
beta1=self.config.trainer.optimizer_adam_beta1,
beta2=self.config.trainer.optimizer_adam_beta2,
)
with tf.control_dependencies(self.update_ops_gen):
self.gen_op = self.gen_optimizer.minimize(
self.loss_generator,
global_step=self.global_step_tensor,
var_list=self.gvars)
with tf.control_dependencies(self.update_ops_enc):
self.enc_op = self.enc_optimizer.minimize(self.loss_encoder,
var_list=self.evars)
with tf.control_dependencies(self.update_ops_dis_xz):
self.dis_op_xz = self.disc_optimizer.minimize(
self.dis_loss_xz, var_list=self.dxzvars)
with tf.control_dependencies(self.update_ops_dis_xx):
self.dis_op_xx = self.disc_optimizer.minimize(
self.dis_loss_xx, var_list=self.dxxvars)
with tf.control_dependencies(self.update_ops_dis_zz):
self.dis_op_zz = self.disc_optimizer.minimize(
self.dis_loss_zz, var_list=self.dzzvars)
# Exponential Moving Average for inference
def train_op_with_ema_dependency(vars, op):
ema = tf.train.ExponentialMovingAverage(
decay=self.config.trainer.ema_decay)
maintain_averages_op = ema.apply(vars)
with tf.control_dependencies([op]):
train_op = tf.group(maintain_averages_op)
return train_op, ema
self.train_gen_op, self.gen_ema = train_op_with_ema_dependency(
self.gvars, self.gen_op)
self.train_enc_op, self.enc_ema = train_op_with_ema_dependency(
self.evars, self.enc_op)
self.train_dis_op_xz, self.xz_ema = train_op_with_ema_dependency(
self.dxzvars, self.dis_op_xz)
self.train_dis_op_xx, self.xx_ema = train_op_with_ema_dependency(
self.dxxvars, self.dis_op_xx)
self.train_dis_op_zz, self.zz_ema = train_op_with_ema_dependency(
self.dzzvars, self.dis_op_zz)
with tf.variable_scope("ALAD"):
with tf.variable_scope("Encoder_Model"):
self.z_gen_ema = self.encoder(
self.image_tensor,
getter=sn.get_getter(self.enc_ema),
do_spectral_norm=self.config.trainer.spectral_norm,
)
with tf.variable_scope("Generator_Model"):
self.rec_x_ema = self.generator(self.z_gen_ema,
getter=sn.get_getter(
self.gen_ema))
self.x_gen_ema = self.generator(self.noise_tensor,
getter=sn.get_getter(
self.gen_ema))
with tf.variable_scope("Discriminator_Model_XX"):
l_encoder_emaxx, inter_layer_inp_emaxx = self.discriminator_xx(
self.image_tensor,
self.image_tensor,
getter=sn.get_getter(self.xx_ema),
do_spectral_norm=self.config.trainer.spectral_norm,
)
l_generator_emaxx, inter_layer_rct_emaxx = self.discriminator_xx(
self.image_tensor,
self.rec_x_ema,
getter=sn.get_getter(self.xx_ema),
do_spectral_norm=self.config.trainer.spectral_norm,
)
with tf.name_scope("Testing"):
with tf.variable_scope("Scores"):
score_ch = tf.nn.sigmoid_cross_entropy_with_logits(
labels=tf.ones_like(l_generator_emaxx),
logits=l_generator_emaxx)
self.score_ch = tf.squeeze(score_ch)
rec = self.image_tensor - self.rec_x_ema
rec = tf.layers.Flatten()(rec)
score_l1 = tf.norm(rec,
ord=1,
axis=1,
keepdims=False,
name="d_loss")
self.score_l1 = tf.squeeze(score_l1)
rec = self.image_tensor - self.rec_x_ema
rec = tf.layers.Flatten()(rec)
score_l2 = tf.norm(rec,
ord=2,
axis=1,
keepdims=False,
name="d_loss")
self.score_l2 = tf.squeeze(score_l2)
inter_layer_inp, inter_layer_rct = (inter_layer_inp_emaxx,
inter_layer_rct_emaxx)
fm = inter_layer_inp - inter_layer_rct
fm = tf.layers.Flatten()(fm)
score_fm = tf.norm(fm,
ord=self.config.trainer.degree,
axis=1,
keepdims=False,
name="d_loss")
self.score_fm = tf.squeeze(score_fm)
if self.config.trainer.enable_early_stop:
self.rec_error_valid = tf.reduce_mean(score_fm)
########################################################################
# TENSORBOARD
########################################################################
if self.config.log.enable_summary:
with tf.name_scope("train_summary"):
with tf.name_scope("dis_summary"):
tf.summary.scalar("loss_discriminator",
self.loss_discriminator, ["dis"])
tf.summary.scalar("loss_dis_encoder", loss_dis_enc,
["dis"])
tf.summary.scalar("loss_dis_gen", loss_dis_gen, ["dis"])
tf.summary.scalar("loss_dis_xz", self.dis_loss_xz, ["dis"])
tf.summary.scalar("loss_dis_xx", self.dis_loss_xx, ["dis"])
if self.config.trainer.allow_zz:
tf.summary.scalar("loss_dis_zz", self.dis_loss_zz,
["dis"])
with tf.name_scope("gen_summary"):
tf.summary.scalar("loss_generator", self.loss_generator,
["gen"])
tf.summary.scalar("loss_encoder", self.loss_encoder,
["gen"])
tf.summary.scalar("loss_encgen_dxx", cost_x, ["gen"])
if self.config.trainer.allow_zz:
tf.summary.scalar("loss_encgen_dzz", cost_z, ["gen"])
with tf.name_scope("img_summary"):
heatmap_pl_latent = tf.placeholder(
tf.float32,
shape=(1, 480, 640, 3),
name="heatmap_pl_latent")
self.sum_op_latent = tf.summary.image(
"heatmap_latent", heatmap_pl_latent)
with tf.name_scope("image_summary"):
tf.summary.image("rec_img", self.rec_img, 3, ["image"])
tf.summary.image("input_image", self.image_tensor, 3,
["image"])
tf.summary.image("image_gen", self.img_gen, 3, ["image"])
tf.summary.image("input_image", self.image_tensor, 3,
["image_2"])
tf.summary.image("rec_x_ema", self.rec_x_ema, 3,
["image_2"])
if self.config.trainer.enable_early_stop:
with tf.name_scope("validation_summary"):
tf.summary.scalar("valid", self.rec_error_valid, ["v"])
self.sum_op_dis = tf.summary.merge_all("dis")
self.sum_op_gen = tf.summary.merge_all("gen")
self.sum_op = tf.summary.merge([self.sum_op_dis, self.sum_op_gen])
self.sum_op_im = tf.summary.merge_all("image")
self.sum_op_im_test = tf.summary.merge_all("image_2")
self.sum_op_valid = tf.summary.merge_all("v")
def build_model(self):
# Initializations
# Kernel initialization for the convolutions
if self.config.trainer.init_type == "normal":
self.init_kernel = tf.random_normal_initializer(mean=0.0,
stddev=0.02)
elif self.config.trainer.init_type == "xavier":
self.init_kernel = tf.contrib.layers.xavier_initializer(
uniform=False, seed=None, dtype=tf.float32)
# Placeholders
self.is_training_gen = tf.placeholder(tf.bool)
self.is_training_dis = tf.placeholder(tf.bool)
self.is_training_enc = tf.placeholder(tf.bool)
self.image_input = tf.placeholder(tf.float32,
shape=[None] +
self.config.trainer.image_dims,
name="x")
self.noise_tensor = tf.placeholder(
tf.float32,
shape=[None, self.config.trainer.noise_dim],
name="noise")
# Build Training Graph
self.logger.info("Building training graph...")
with tf.variable_scope("EncEBGAN"):
with tf.variable_scope("Generator_Model"):
self.image_gen = self.generator(self.noise_tensor)
with tf.variable_scope("Discriminator_Model"):
self.embedding_real, self.decoded_real = self.discriminator(
self.image_input,
do_spectral_norm=self.config.trainer.do_spectral_norm)
self.embedding_fake, self.decoded_fake = self.discriminator(
self.image_gen,
do_spectral_norm=self.config.trainer.do_spectral_norm)
with tf.variable_scope("Encoder_Model"):
self.image_encoded = self.encoder(self.image_input)
with tf.variable_scope("Generator_Model"):
self.image_gen_enc = self.generator(self.image_encoded)
with tf.variable_scope("Discriminator_Model"):
self.embedding_enc_fake, self.decoded_enc_fake = self.discriminator(
self.image_gen_enc,
do_spectral_norm=self.config.trainer.do_spectral_norm)
self.embedding_enc_real, self.decoded_enc_real = self.discriminator(
self.image_input,
do_spectral_norm=self.config.trainer.do_spectral_norm)
# Loss functions
with tf.name_scope("Loss_Functions"):
with tf.name_scope("Generator_Discriminator"):
# Discriminator Loss
if self.config.trainer.mse_mode == "norm":
self.disc_loss_real = tf.reduce_mean(
self.mse_loss(self.decoded_real,
self.image_input,
mode="norm"))
self.disc_loss_fake = tf.reduce_mean(
self.mse_loss(self.decoded_fake,
self.image_gen,
mode="norm"))
elif self.config.trainer.mse_mode == "mse":
self.disc_loss_real = self.mse_loss(self.decoded_real,
self.image_input,
mode="mse")
self.disc_loss_fake = self.mse_loss(self.decoded_fake,
self.image_gen,
mode="mse")
self.loss_discriminator = (tf.math.maximum(
self.config.trainer.disc_margin - self.disc_loss_fake, 0) +
self.disc_loss_real)
# Generator Loss
pt_loss = 0
if self.config.trainer.pullaway:
pt_loss = self.pullaway_loss(self.embedding_fake)
self.loss_generator = self.disc_loss_fake + self.config.trainer.pt_weight * pt_loss
with tf.name_scope("Encoder"):
if self.config.trainer.mse_mode == "norm":
self.loss_enc_rec = tf.reduce_mean(
self.mse_loss(self.image_gen_enc,
self.image_input,
mode="norm"))
self.loss_enc_f = tf.reduce_mean(
self.mse_loss(self.embedding_enc_real,
self.embedding_enc_fake,
mode="norm"))
elif self.config.trainer.mse_mode == "mse":
self.loss_enc_rec = tf.reduce_mean(
self.mse_loss(self.image_gen_enc,
self.image_input,
mode="mse"))
self.loss_enc_f = tf.reduce_mean(
self.mse_loss(self.embedding_enc_real,
self.embedding_enc_fake,
mode="mse"))
self.loss_encoder = (
self.loss_enc_rec +
self.config.trainer.encoder_f_factor * self.loss_enc_f)
# Optimizers
with tf.name_scope("Optimizers"):
self.generator_optimizer = tf.train.AdamOptimizer(
self.config.trainer.standard_lr_gen,
beta1=self.config.trainer.optimizer_adam_beta1,
beta2=self.config.trainer.optimizer_adam_beta2,
)
self.encoder_optimizer = tf.train.AdamOptimizer(
self.config.trainer.standard_lr_enc,
beta1=self.config.trainer.optimizer_adam_beta1,
beta2=self.config.trainer.optimizer_adam_beta2,
)
self.discriminator_optimizer = tf.train.AdamOptimizer(
self.config.trainer.standard_lr_dis,
beta1=self.config.trainer.optimizer_adam_beta1,
beta2=self.config.trainer.optimizer_adam_beta2,
)
# Collect all the variables
all_variables = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES)
# Generator Network Variables
self.generator_vars = [
v for v in all_variables
if v.name.startswith("EncEBGAN/Generator_Model")
]
# Discriminator Network Variables
self.discriminator_vars = [
v for v in all_variables
if v.name.startswith("EncEBGAN/Discriminator_Model")
]
# Discriminator Network Variables
self.encoder_vars = [
v for v in all_variables
if v.name.startswith("EncEBGAN/Encoder_Model")
]
# Generator Network Operations
self.gen_update_ops = tf.get_collection(
tf.GraphKeys.UPDATE_OPS, scope="EncEBGAN/Generator_Model")
# Discriminator Network Operations
self.disc_update_ops = tf.get_collection(
tf.GraphKeys.UPDATE_OPS, scope="EncEBGAN/Discriminator_Model")
self.enc_update_ops = tf.get_collection(
tf.GraphKeys.UPDATE_OPS, scope="EncEBGAN/Encoder_Model")
with tf.control_dependencies(self.gen_update_ops):
self.gen_op = self.generator_optimizer.minimize(
self.loss_generator,
var_list=self.generator_vars,
global_step=self.global_step_tensor,
)
with tf.control_dependencies(self.disc_update_ops):
self.disc_op = self.discriminator_optimizer.minimize(
self.loss_discriminator, var_list=self.discriminator_vars)
with tf.control_dependencies(self.enc_update_ops):
self.enc_op = self.encoder_optimizer.minimize(
self.loss_encoder,
var_list=self.encoder_vars,
global_step=self.global_step_tensor,
)
# Exponential Moving Average for Estimation
self.dis_ema = tf.train.ExponentialMovingAverage(
decay=self.config.trainer.ema_decay)
maintain_averages_op_dis = self.dis_ema.apply(
self.discriminator_vars)
self.gen_ema = tf.train.ExponentialMovingAverage(
decay=self.config.trainer.ema_decay)
maintain_averages_op_gen = self.gen_ema.apply(self.generator_vars)
self.enc_ema = tf.train.ExponentialMovingAverage(
decay=self.config.trainer.ema_decay)
maintain_averages_op_enc = self.enc_ema.apply(self.encoder_vars)
with tf.control_dependencies([self.disc_op]):
self.train_dis_op = tf.group(maintain_averages_op_dis)
with tf.control_dependencies([self.gen_op]):
self.train_gen_op = tf.group(maintain_averages_op_gen)
with tf.control_dependencies([self.enc_op]):
self.train_enc_op = tf.group(maintain_averages_op_enc)
# Build Test Graph
self.logger.info("Building Testing Graph...")
with tf.variable_scope("EncEBGAN"):
with tf.variable_scope("Discriminator_Model"):
self.embedding_q_ema, self.decoded_q_ema = self.discriminator(
self.image_input,
getter=get_getter(self.dis_ema),
do_spectral_norm=self.config.trainer.do_spectral_norm,
)
with tf.variable_scope("Generator_Model"):
self.image_gen_ema = self.generator(self.embedding_q_ema,
getter=get_getter(
self.gen_ema))
with tf.variable_scope("Discriminator_Model"):
self.embedding_rec_ema, self.decoded_rec_ema = self.discriminator(
self.image_gen_ema,
getter=get_getter(self.dis_ema),
do_spectral_norm=self.config.trainer.do_spectral_norm,
)
with tf.variable_scope("Encoder_Model"):
self.image_encoded_ema = self.encoder(self.image_input,
getter=get_getter(
self.enc_ema))
with tf.variable_scope("Generator_Model"):
self.image_gen_enc_ema = self.generator(self.image_encoded_ema,
getter=get_getter(
self.gen_ema))
with tf.variable_scope("Discriminator_Model"):
self.embedding_enc_fake_ema, self.decoded_enc_fake_ema = self.discriminator(
self.image_gen_enc_ema,
getter=get_getter(self.dis_ema),
do_spectral_norm=self.config.trainer.do_spectral_norm,
)
self.embedding_enc_real_ema, self.decoded_enc_real_ema = self.discriminator(
self.image_input,
getter=get_getter(self.dis_ema),
do_spectral_norm=self.config.trainer.do_spectral_norm,
)
with tf.name_scope("Testing"):
with tf.name_scope("Image_Based"):
delta = self.image_input - self.image_gen_enc_ema
self.rec_residual = -delta
delta_flat = tf.layers.Flatten()(delta)
img_score_l1 = tf.norm(delta_flat,
ord=2,
axis=1,
keepdims=False,
name="img_loss__1")
self.img_score_l1 = tf.squeeze(img_score_l1)
delta = self.decoded_enc_fake_ema - self.decoded_enc_real_ema
delta_flat = tf.layers.Flatten()(delta)
img_score_l2 = tf.norm(delta_flat,
ord=2,
axis=1,
keepdims=False,
name="img_loss__2")
self.img_score_l2 = tf.squeeze(img_score_l2)
self.score_comb = (
(1 - self.config.trainer.feature_match_weight) *
self.img_score_l1 +
self.config.trainer.feature_match_weight *
self.img_score_l2)
with tf.name_scope("Noise_Based"):
delta = self.image_encoded_ema - self.embedding_enc_fake_ema
delta_flat = tf.layers.Flatten()(delta)
z_score_l1 = tf.norm(delta_flat,
ord=2,
axis=1,
keepdims=False,
name="z_loss_1")
self.z_score_l1 = tf.squeeze(z_score_l1)
delta = self.embedding_enc_real_ema - self.embedding_enc_fake_ema
delta_flat = tf.layers.Flatten()(delta)
z_score_l2 = tf.norm(delta_flat,
ord=2,
axis=1,
keepdims=False,
name="z_loss_2")
self.z_score_l2 = tf.squeeze(z_score_l2)
self.score_comb_2 = (
(1 - self.config.trainer.feature_match_weight) *
self.z_score_l1 +
self.config.trainer.feature_match_weight * self.z_score_l2)
# Tensorboard
if self.config.log.enable_summary:
with tf.name_scope("train_summary"):
with tf.name_scope("dis_summary"):
tf.summary.scalar("loss_disc", self.loss_discriminator,
["dis"])
tf.summary.scalar("loss_disc_real", self.disc_loss_real,
["dis"])
tf.summary.scalar("loss_disc_fake", self.disc_loss_fake,
["dis"])
with tf.name_scope("gen_summary"):
tf.summary.scalar("loss_generator", self.loss_generator,
["gen"])
with tf.name_scope("enc_summary"):
tf.summary.scalar("loss_encoder", self.loss_encoder,
["enc"])
with tf.name_scope("img_summary"):
tf.summary.image("input_image", self.image_input, 1,
["img_1"])
tf.summary.image("reconstructed", self.image_gen, 1,
["img_1"])
tf.summary.image("input_enc", self.image_input, 1,
["img_2"])
tf.summary.image("reconstructed", self.image_gen_enc, 1,
["img_2"])
tf.summary.image("input_image", self.image_input, 1,
["test"])
tf.summary.image("reconstructed", self.image_gen_enc_ema,
1, ["test"])
tf.summary.image("residual", self.rec_residual, 1,
["test"])
self.sum_op_dis = tf.summary.merge_all("dis")
self.sum_op_gen = tf.summary.merge_all("gen")
self.sum_op_enc = tf.summary.merge_all("enc")
self.sum_op_im_1 = tf.summary.merge_all("img_1")
self.sum_op_im_2 = tf.summary.merge_all("img_2")
self.sum_op_im_test = tf.summary.merge_all("test")
self.sum_op = tf.summary.merge([self.sum_op_dis, self.sum_op_gen])
def build_model(self):
# Kernel initialization for the convolutions
if self.config.trainer.init_type == "normal":
self.init_kernel = tf.random_normal_initializer(mean=0.0, stddev=0.02)
elif self.config.trainer.init_type == "xavier":
self.init_kernel = tf.contrib.layers.xavier_initializer(
uniform=False, seed=None, dtype=tf.float32
)
# Placeholders
self.is_training_gen = tf.placeholder(tf.bool)
self.is_training_dis = tf.placeholder(tf.bool)
self.is_training_enc = tf.placeholder(tf.bool)
self.image_input = tf.placeholder(
tf.float32, shape=[None] + self.config.trainer.image_dims, name="x"
)
self.noise_tensor = tf.placeholder(
tf.float32, shape=[None, self.config.trainer.noise_dim], name="noise"
)
self.true_labels = tf.placeholder(dtype=tf.float32, shape=[None, 1], name="true_labels")
self.generated_labels = tf.placeholder(dtype=tf.float32, shape=[None, 1], name="gen_labels")
self.real_noise = tf.placeholder(
dtype=tf.float32, shape=[None] + self.config.trainer.image_dims, name="real_noise"
)
self.fake_noise = tf.placeholder(
dtype=tf.float32, shape=[None] + self.config.trainer.image_dims, name="fake_noise"
)
self.logger.info("Building training graph...")
with tf.variable_scope("FAnogan"):
# Generator and Discriminator Training
with tf.variable_scope("Generator_Model"):
self.image_gen = self.generator(self.noise_tensor) + self.fake_noise
with tf.variable_scope("Discriminator_Model"):
self.disc_real, self.disc_f_real = self.discriminator(
self.image_input + self.real_noise
)
self.disc_fake, self.disc_f_fake = self.discriminator(self.image_gen)
# Encoder Training
with tf.variable_scope("Encoder_Model"):
# ZIZ Architecture
self.encoded_gen_noise = self.encoder(self.image_gen)
# IZI Architecture
self.encoded_img = self.encoder(self.image_input)
with tf.variable_scope("Generator_Model"):
self.gen_enc_img = self.generator(self.encoded_img)
with tf.variable_scope("Discriminator_Model"):
# IZI Training
self.disc_real_izi, self.disc_f_real_izi = self.discriminator(self.image_input)
self.disc_fake_izi, self.disc_f_fake_izi = self.discriminator(self.gen_enc_img)
with tf.name_scope("Loss_Funcions"):
with tf.name_scope("Encoder"):
if self.config.trainer.encoder_training_mode == "ziz":
self.loss_encoder = tf.reduce_mean(
self.mse_loss(
self.encoded_gen_noise,
self.noise_tensor,
mode=self.config.trainer.encoder_loss_mode,
)
* (1.0 / self.config.trainer.noise_dim)
)
elif self.config.trainer.encoder_training_mode == "izi":
self.izi_reconstruction = self.mse_loss(
self.image_input,
self.gen_enc_img,
mode=self.config.trainer.encoder_loss_mode,
) * (
1.0
/ (self.config.data_loader.image_size * self.config.data_loader.image_size)
)
self.loss_encoder = tf.reduce_mean(self.izi_reconstruction)
elif self.config.trainer.encoder_training_mode == "izi_f":
self.izi_reconstruction = self.mse_loss(
self.image_input,
self.gen_enc_img,
mode=self.config.trainer.encoder_loss_mode,
) * (
1.0
/ (self.config.data_loader.image_size * self.config.data_loader.image_size)
)
self.izi_disc = self.mse_loss(
self.disc_f_real_izi,
self.disc_f_fake_izi,
mode=self.config.trainer.encoder_loss_mode,
) * (
1.0
* self.config.trainer.kappa_weight_factor
/ self.config.trainer.feature_layer_dim
)
self.loss_encoder = tf.reduce_mean(self.izi_reconstruction + self.izi_disc)
with tf.name_scope("Discriminator_Generator"):
if self.config.trainer.mode == "standard":
self.loss_disc_real = tf.reduce_mean(
tf.nn.sigmoid_cross_entropy_with_logits(
labels=self.true_labels, logits=self.disc_real
)
)
self.loss_disc_fake = tf.reduce_mean(
tf.nn.sigmoid_cross_entropy_with_logits(
labels=self.generated_labels, logits=self.disc_fake
)
)
self.loss_discriminator = self.loss_disc_real + self.loss_disc_fake
# Flip the weigths for the encoder and generator
if self.config.trainer.flip_labels:
labels_gen = tf.zeros_like(self.disc_fake)
else:
labels_gen = tf.ones_like(self.disc_fake)
# Generator
self.loss_generator_ce = tf.reduce_mean(
tf.nn.sigmoid_cross_entropy_with_logits(
labels=labels_gen, logits=self.disc_fake
)
)
delta = self.disc_f_fake - self.disc_f_real
delta = tf.layers.Flatten()(delta)
self.loss_generator_fm = tf.reduce_mean(
tf.norm(delta, ord=2, axis=1, keepdims=False)
)
self.loss_generator = (
self.loss_generator_ce
+ self.config.trainer.feature_match_weight * self.loss_generator_fm
)
elif self.config.trainer.mode == "wgan":
self.loss_d_fake = -tf.reduce_mean(self.disc_fake)
self.loss_d_real = -tf.reduce_mean(self.disc_real)
self.loss_discriminator = -self.loss_d_fake + self.loss_d_real
self.loss_generator = -tf.reduce_mean(self.disc_fake)
# Weight Clipping and Encoder Part
elif self.config.trainer.mode == "wgan_gp":
self.loss_generator = -tf.reduce_mean(self.disc_fake)
self.loss_d_fake = -tf.reduce_mean(self.disc_fake)
self.loss_d_real = -tf.reduce_mean(self.disc_real)
self.loss_discriminator = -self.loss_d_fake - self.loss_d_real
alpha_x = tf.random_uniform(
shape=[self.config.data_loader.batch_size] + self.config.trainer.image_dims,
minval=0.0,
maxval=1.0,
)
differences_x = self.image_gen - self.image_input
interpolates_x = self.image_input + (alpha_x * differences_x)
gradients = tf.gradients(self.discriminator(interpolates_x), [interpolates_x])[
0
]
slopes = tf.sqrt(tf.reduce_sum(tf.square(gradients), reduction_indices=[1]))
gradient_penalty = tf.reduce_mean((slopes - 1.0) ** 2)
self.loss_discriminator += self.config.trainer.wgan_gp_lambda * gradient_penalty
with tf.name_scope("Optimizations"):
if self.config.trainer.mode == "standard":
# Build the optimizers
self.generator_optimizer = tf.train.AdamOptimizer(
self.config.trainer.standard_lr,
beta1=self.config.trainer.optimizer_adam_beta1,
beta2=self.config.trainer.optimizer_adam_beta2,
)
self.discriminator_optimizer = tf.train.AdamOptimizer(
self.config.trainer.standard_lr_disc,
beta1=self.config.trainer.optimizer_adam_beta1,
beta2=self.config.trainer.optimizer_adam_beta2,
)
self.encoder_optimizer = tf.train.AdamOptimizer(
self.config.trainer.standard_lr,
beta1=self.config.trainer.optimizer_adam_beta1,
beta2=self.config.trainer.optimizer_adam_beta2,
)
elif self.config.trainer.mode == "wgan":
# Build the optimizers
self.generator_optimizer = tf.train.RMSPropOptimizer(self.config.trainer.wgan_lr)
self.discriminator_optimizer = tf.train.RMSPropOptimizer(
self.config.trainer.standard_lr_disc
)
self.encoder_optimizer = tf.train.AdamOptimizer(
self.config.trainer.wgan_lr,
beta1=self.config.trainer.optimizer_adam_beta1,
beta2=self.config.trainer.optimizer_adam_beta2,
)
elif self.config.trainer.mode == "wgan_gp":
# Build the optimizers
self.generator_optimizer = tf.train.AdamOptimizer(
self.config.trainer.wgan_gp_lr, beta1=0.0, beta2=0.9
)
self.discriminator_optimizer = tf.train.AdamOptimizer(
self.config.trainer.standard_lr_disc, beta1=0.0, beta2=0.9
)
self.encoder_optimizer = tf.train.AdamOptimizer(
self.config.trainer.wgan_gp_lr, beta1=0.0, beta2=0.9
)
# Collect all the variables
all_variables = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES)
# Generator Network Variables
self.generator_vars = [
v for v in all_variables if v.name.startswith("FAnogan/Generator_Model")
]
# Discriminator Network Variables
self.discriminator_vars = [
v for v in all_variables if v.name.startswith("FAnogan/Discriminator_Model")
]
if self.config.trainer.mode == "wgan":
clip_ops = []
for var in self.discriminator_vars:
clip_bounds = [-0.01, 0.01]
clip_ops.append(
tf.assign(var, tf.clip_by_value(var, clip_bounds[0], clip_bounds[1]))
)
self.clip_disc_weights = tf.group(*clip_ops)
# Encoder Network Variables
self.encoder_vars = [
v for v in all_variables if v.name.startswith("FAnogan/Encoder_Model")
]
# Create Training Operations
# Generator Network Operations
self.gen_update_ops = tf.get_collection(
tf.GraphKeys.UPDATE_OPS, scope="FAnogan/Generator_Model"
)
# Discriminator Network Operations
self.disc_update_ops = tf.get_collection(
tf.GraphKeys.UPDATE_OPS, scope="FAnogan/Discriminator_Model"
)
# Encoder Network Operations
self.enc_update_ops = tf.get_collection(
tf.GraphKeys.UPDATE_OPS, scope="FAnogan/Encoder_Model"
)
# Initialization of Optimizers
with tf.control_dependencies(self.gen_update_ops):
self.gen_op = self.generator_optimizer.minimize(
self.loss_generator,
var_list=self.generator_vars,
global_step=self.global_step_tensor,
)
with tf.control_dependencies(self.disc_update_ops):
self.disc_op = self.discriminator_optimizer.minimize(
self.loss_discriminator, var_list=self.discriminator_vars
)
with tf.control_dependencies(self.enc_update_ops):
self.enc_op = self.encoder_optimizer.minimize(
self.loss_encoder, var_list=self.encoder_vars
)
# Exponential Moving Average for Estimation
self.dis_ema = tf.train.ExponentialMovingAverage(decay=self.config.trainer.ema_decay)
maintain_averages_op_dis = self.dis_ema.apply(self.discriminator_vars)
self.gen_ema = tf.train.ExponentialMovingAverage(decay=self.config.trainer.ema_decay)
maintain_averages_op_gen = self.gen_ema.apply(self.generator_vars)
self.enc_ema = tf.train.ExponentialMovingAverage(decay=self.config.trainer.ema_decay)
maintain_averages_op_enc = self.enc_ema.apply(self.encoder_vars)
with tf.control_dependencies([self.disc_op]):
self.train_dis_op = tf.group(maintain_averages_op_dis)
with tf.control_dependencies([self.gen_op]):
self.train_gen_op = tf.group(maintain_averages_op_gen)
with tf.control_dependencies([self.enc_op]):
self.train_enc_op = tf.group(maintain_averages_op_enc)
self.logger.info("Building Testing Graph...")
with tf.variable_scope("FAnogan"):
# Generator and Discriminator Training
with tf.variable_scope("Generator_Model"):
self.image_gen_ema = self.generator(
self.noise_tensor, getter=get_getter(self.gen_ema)
)
with tf.variable_scope("Discriminator_Model"):
self.disc_real_ema, self.disc_f_real_ema = self.discriminator(
self.image_input, getter=get_getter(self.dis_ema)
)
self.disc_fake_ema, self.disc_f_fake_ema = self.discriminator(
self.image_gen_ema, getter=get_getter(self.dis_ema)
)
# Encoder Training
with tf.variable_scope("Encoder_Model"):
# ZIZ Architecture
self.encoded_gen_noise_ema = self.encoder(
self.image_gen_ema, getter=get_getter(self.enc_ema)
)
# IZI Architecture
self.encoded_img_ema = self.encoder(
self.image_input, getter=get_getter(self.enc_ema)
)
with tf.variable_scope("Generator_Model"):
self.gen_enc_img_ema = self.generator(
self.encoded_img_ema, getter=get_getter(self.gen_ema)
)
with tf.variable_scope("Discriminator_Model"):
# IZI Training
self.disc_real_izi_ema, self.disc_f_real_izi_ema = self.discriminator(
self.image_input, getter=get_getter(self.dis_ema)
)
self.disc_fake_izi_ema, self.disc_f_fake_izi_ema = self.discriminator(
self.gen_enc_img_ema, getter=get_getter(self.dis_ema)
)
with tf.name_scope("Testing"):
with tf.name_scope("izi_f_loss"):
self.score_reconstruction = self.mse_loss(
self.image_input, self.gen_enc_img_ema
) * (
1.0 / (self.config.data_loader.image_size * self.config.data_loader.image_size)
)
self.score_disc = self.mse_loss(
self.disc_f_real_izi_ema, self.disc_f_fake_izi_ema
) * (
1.0
* self.config.trainer.kappa_weight_factor
/ self.config.trainer.feature_layer_dim
)
self.izi_f_score = self.score_reconstruction + self.score_disc
with tf.name_scope("ziz_loss"):
self.score_reconstruction = self.mse_loss(
self.image_input, self.gen_enc_img_ema
) * (
1.0 / (self.config.data_loader.image_size * self.config.data_loader.image_size)
)
self.ziz_score = self.score_reconstruction
if self.config.trainer.enable_early_stop:
self.rec_error_valid = tf.reduce_mean(self.izi_f_score)
if self.config.log.enable_summary:
with tf.name_scope("Summary"):
with tf.name_scope("Disc_Summary"):
tf.summary.scalar("loss_discriminator", self.loss_discriminator, ["dis"])
if self.config.trainer.mode == "standard":
tf.summary.scalar("loss_dis_real", self.loss_disc_real, ["dis"])
tf.summary.scalar("loss_dis_fake", self.loss_disc_fake, ["dis"])
with tf.name_scope("Gen_Summary"):
tf.summary.scalar("loss_generator", self.loss_generator, ["gen"])
if self.config.trainer.mode == "standard":
tf.summary.scalar("loss_generator_ce", self.loss_generator_ce, ["gen"])
tf.summary.scalar("loss_generator_fm", self.loss_generator_fm, ["gen"])
tf.summary.scalar("loss_encoder", self.loss_encoder, ["enc"])
with tf.name_scope("Image_Summary"):
tf.summary.image("reconstruct", self.image_gen, 3, ["image_1"])
tf.summary.image("input_images", self.image_input, 3, ["image_1"])
tf.summary.image("gen_enc_img", self.gen_enc_img, 3, ["image_2"])
tf.summary.image("input_image_2", self.image_input, 3, ["image_2"])
if self.config.trainer.enable_early_stop:
with tf.name_scope("validation_summary"):
tf.summary.scalar("valid", self.rec_error_valid, ["v"])
self.sum_op_dis = tf.summary.merge_all("dis")
self.sum_op_gen = tf.summary.merge_all("gen")
self.sum_op_enc = tf.summary.merge_all("enc")
self.sum_op_im_1 = tf.summary.merge_all("image_1")
self.sum_op_im_2 = tf.summary.merge_all("image_2")
self.sum_op_valid = tf.summary.merge_all("v")
def build_model(self):
# Kernel initialization for the convolutions
self.init_kernel = tf.random_normal_initializer(mean=0.0, stddev=0.02)
# Placeholders
self.is_training = tf.placeholder(tf.bool)
self.image_input = tf.placeholder(tf.float32,
shape=[None] +
self.config.trainer.image_dims,
name="x")
self.true_labels = tf.placeholder(dtype=tf.float32,
shape=[None, 1],
name="true_labels")
self.generated_labels = tf.placeholder(dtype=tf.float32,
shape=[None, 1],
name="gen_labels")
self.logger.info("Building training graph...")
with tf.variable_scope("GANomaly"):
with tf.variable_scope("Generator_Model"):
self.noise_gen, self.img_rec, self.noise_rec = self.generator(
self.image_input)
with tf.variable_scope("Discriminator_Model"):
l_real, inter_layer_inp = self.discriminator(self.image_input)
l_fake, inter_layer_rct = self.discriminator(self.img_rec)
with tf.name_scope("Loss_Functions"):
# Discriminator
self.loss_dis_real = tf.reduce_mean(
tf.nn.sigmoid_cross_entropy_with_logits(
labels=self.true_labels, logits=l_real))
self.loss_dis_fake = tf.reduce_mean(
tf.nn.sigmoid_cross_entropy_with_logits(
labels=self.generated_labels, logits=l_fake))
# Feature matching part
fm = inter_layer_inp - inter_layer_rct
fm = tf.layers.Flatten()(fm)
self.feature_match = tf.reduce_mean(
tf.norm(fm, ord=2, axis=1, keepdims=False))
self.loss_discriminator = (self.loss_dis_fake +
self.loss_dis_real + self.feature_match
if self.config.trainer.loss_method
== "fm" else self.loss_dis_fake +
self.loss_dis_real)
# Generator
# Adversarial Loss
if self.config.trainer.flip_labels:
labels = tf.zeros_like(l_fake)
else:
labels = tf.ones_like(l_real)
self.gen_loss_ce = tf.reduce_mean(
tf.nn.sigmoid_cross_entropy_with_logits(labels=labels,
logits=l_fake))
# Contextual Loss
l1_norm = self.image_input - self.img_rec
l1_norm = tf.layers.Flatten()(l1_norm)
self.gen_loss_con = tf.reduce_mean(
tf.norm(l1_norm, ord=1, axis=1, keepdims=False))
# Encoder Loss
l2_norm = self.noise_gen - self.noise_rec
l2_norm = tf.layers.Flatten()(l2_norm)
self.gen_loss_enc = tf.reduce_mean(
tf.norm(l2_norm, ord=2, axis=1, keepdims=False))
self.gen_loss_total = (
self.config.trainer.weight_adv * self.gen_loss_ce +
self.config.trainer.weight_cont * self.gen_loss_con +
self.config.trainer.weight_enc * self.gen_loss_enc)
with tf.name_scope("Optimizers"):
# Build the optimizers
self.generator_optimizer = tf.train.AdamOptimizer(
self.config.trainer.generator_l_rate,
beta1=self.config.trainer.optimizer_adam_beta1,
beta2=self.config.trainer.optimizer_adam_beta2,
)
self.discriminator_optimizer = tf.train.AdamOptimizer(
self.config.trainer.discriminator_l_rate,
beta1=self.config.trainer.optimizer_adam_beta1,
beta2=self.config.trainer.optimizer_adam_beta2,
)
# Collect all the variables
all_variables = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES)
# Generator Network Variables
self.generator_vars = [
v for v in all_variables
if v.name.startswith("GANomaly/Generator_Model")
]
# Discriminator Network Variables
self.discriminator_vars = [
v for v in all_variables
if v.name.startswith("GANomaly/Discriminator_Model")
]
# Create Training Operations
# Generator Network Operations
self.gen_update_ops = tf.get_collection(
tf.GraphKeys.UPDATE_OPS, scope="GANomaly/Generator_Model")
# Discriminator Network Operations
self.disc_update_ops = tf.get_collection(
tf.GraphKeys.UPDATE_OPS, scope="GANomaly/Discriminator_Model")
# Initialization of Optimizers
with tf.control_dependencies(self.gen_update_ops):
self.gen_op = self.generator_optimizer.minimize(
self.gen_loss_total,
global_step=self.global_step_tensor,
var_list=self.generator_vars,
)
with tf.control_dependencies(self.disc_update_ops):
self.disc_op = self.discriminator_optimizer.minimize(
self.loss_discriminator, var_list=self.discriminator_vars)
# Exponential Moving Average for Estimation
self.dis_ema = tf.train.ExponentialMovingAverage(
decay=self.config.trainer.ema_decay)
maintain_averages_op_dis = self.dis_ema.apply(
self.discriminator_vars)
self.gen_ema = tf.train.ExponentialMovingAverage(
decay=self.config.trainer.ema_decay)
maintain_averages_op_gen = self.gen_ema.apply(self.generator_vars)
with tf.control_dependencies([self.disc_op]):
self.train_dis_op = tf.group(maintain_averages_op_dis)
with tf.control_dependencies([self.gen_op]):
self.train_gen_op = tf.group(maintain_averages_op_gen)
self.logger.info("Building Testing Graph...")
with tf.variable_scope("GANomaly"):
with tf.variable_scope("Generator_Model"):
self.noise_gen_ema, self.img_rec_ema, self.noise_rec_ema = self.generator(
self.image_input, getter=get_getter(self.gen_ema))
with tf.variable_scope("Discriminator_model"):
self.l_real_ema, self.inter_layer_inp_ema = self.discriminator(
self.image_input, getter=get_getter(self.dis_ema))
self.l_fake_ema, self.inter_layer_rct_ema = self.discriminator(
self.img_rec_ema, getter=get_getter(self.dis_ema))
with tf.name_scope("Testing"):
with tf.variable_scope("Reconstruction_Loss"):
# | G_E(x) - E(G(x))|1
# Difference between the noise generated from the input image and reconstructed noise
delta = self.noise_gen_ema - self.noise_rec_ema
delta = tf.layers.Flatten()(delta)
self.score = tf.norm(delta, ord=1, axis=1, keepdims=False)
if self.config.trainer.enable_early_stop:
self.rec_error_valid = tf.reduce_mean(self.score)
if self.config.log.enable_summary:
with tf.name_scope("summary"):
with tf.name_scope("disc_summary"):
tf.summary.scalar("loss_discriminator_total",
self.loss_discriminator, ["dis"])
tf.summary.scalar("loss_dis_real", self.loss_dis_real,
["dis"])
tf.summary.scalar("loss_dis_fake", self.loss_dis_fake,
["dis"])
if self.config.trainer.loss_method:
tf.summary.scalar("loss_dis_fm", self.feature_match,
["dis"])
with tf.name_scope("gen_summary"):
tf.summary.scalar("loss_generator_total",
self.gen_loss_total, ["gen"])
tf.summary.scalar("loss_gen_adv", self.gen_loss_ce,
["gen"])
tf.summary.scalar("loss_gen_con", self.gen_loss_con,
["gen"])
tf.summary.scalar("loss_gen_enc", self.gen_loss_enc,
["gen"])
with tf.name_scope("image_summary"):
tf.summary.image("reconstruct", self.img_rec, 3, ["image"])
tf.summary.image("input_images", self.image_input, 3,
["image"])
if self.config.trainer.enable_early_stop:
with tf.name_scope("validation_summary"):
tf.summary.scalar("valid", self.rec_error_valid, ["v"])
self.sum_op_dis = tf.summary.merge_all("dis")
self.sum_op_gen = tf.summary.merge_all("gen")
self.sum_op_im = tf.summary.merge_all("image")
self.sum_op_valid = tf.summary.merge_all("v")
def build_model(self):
# Placeholdersn
if self.config.trainer.init_type == "normal":
self.init_kernel = tf.random_normal_initializer(mean=0.0,
stddev=0.02)
elif self.config.trainer.init_type == "xavier":
self.init_kernel = tf.contrib.layers.xavier_initializer(
uniform=False, seed=None, dtype=tf.float32)
self.is_training = tf.placeholder(tf.bool)
self.image_input = tf.placeholder(tf.float32,
shape=[None] +
self.config.trainer.image_dims,
name="x")
self.noise_tensor = tf.placeholder(
tf.float32,
shape=[None, self.config.trainer.noise_dim],
name="noise")
# Placeholders for the true and fake labels
self.true_labels = tf.placeholder(dtype=tf.float32, shape=[None, 1])
self.generated_labels = tf.placeholder(dtype=tf.float32,
shape=[None, 1])
self.real_noise = tf.placeholder(dtype=tf.float32,
shape=[None] +
self.config.trainer.image_dims,
name="real_noise")
self.fake_noise = tf.placeholder(dtype=tf.float32,
shape=[None] +
self.config.trainer.image_dims,
name="fake_noise")
# Building the Graph
self.logger.info("Building Graph")
with tf.variable_scope("ANOGAN"):
with tf.variable_scope("Generator_Model"):
self.img_gen = self.generator(
self.noise_tensor) + self.fake_noise
# Discriminator
with tf.variable_scope("Discriminator_Model"):
disc_real, inter_layer_real = self.discriminator(
self.image_input + self.real_noise)
disc_fake, inter_layer_fake = self.discriminator(self.img_gen)
# Losses of the training of Generator and Discriminator
with tf.variable_scope("Loss_Functions"):
with tf.name_scope("Discriminator_Loss"):
self.disc_loss_real = tf.reduce_mean(
tf.nn.sigmoid_cross_entropy_with_logits(
labels=self.true_labels, logits=disc_real))
self.disc_loss_fake = tf.reduce_mean(
tf.nn.sigmoid_cross_entropy_with_logits(
labels=self.generated_labels, logits=disc_fake))
self.total_disc_loss = self.disc_loss_real + self.disc_loss_fake
with tf.name_scope("Generator_Loss"):
if self.config.trainer.flip_labels:
labels = tf.zeros_like(disc_fake)
else:
labels = tf.ones_like(disc_fake)
self.gen_loss = tf.reduce_mean(
tf.nn.sigmoid_cross_entropy_with_logits(labels=labels,
logits=disc_fake))
# Build the Optimizers
with tf.variable_scope("Optimizers"):
# Collect all the variables
all_variables = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES)
# Generator Network Variables
self.generator_vars = [
v for v in all_variables
if v.name.startswith("ANOGAN/Generator_Model")
]
# Discriminator Network Variables
self.discriminator_vars = [
v for v in all_variables
if v.name.startswith("ANOGAN/Discriminator_Model")
]
# Create Training Operations
# Generator Network Operations
self.gen_update_ops = tf.get_collection(
tf.GraphKeys.UPDATE_OPS, scope="ANOGAN/Generator_Model")
# Discriminator Network Operations
self.disc_update_ops = tf.get_collection(
tf.GraphKeys.UPDATE_OPS, scope="ANOGAN/Discriminator_Model")
# Initialization of Optimizers
self.generator_optimizer = tf.train.AdamOptimizer(
self.config.trainer.generator_l_rate,
beta1=self.config.trainer.optimizer_adam_beta1,
beta2=self.config.trainer.optimizer_adam_beta2,
)
self.discriminator_optimizer = tf.train.AdamOptimizer(
self.config.trainer.discriminator_l_rate,
beta1=self.config.trainer.optimizer_adam_beta1,
beta2=self.config.trainer.optimizer_adam_beta2,
)
with tf.control_dependencies(self.gen_update_ops):
self.train_gen = self.generator_optimizer.minimize(
self.gen_loss,
global_step=self.global_step_tensor,
var_list=self.generator_vars)
with tf.control_dependencies(self.disc_update_ops):
self.train_disc = self.discriminator_optimizer.minimize(
self.total_disc_loss, var_list=self.discriminator_vars)
def train_op_with_ema_dependency(vars, op):
ema = tf.train.ExponentialMovingAverage(
decay=self.config.trainer.ema_decay)
maintain_averages_op = ema.apply(vars)
with tf.control_dependencies([op]):
train_op = tf.group(maintain_averages_op)
return train_op, ema
self.train_gen_op, self.gen_ema = train_op_with_ema_dependency(
self.generator_vars, self.train_gen)
self.train_dis_op, self.dis_ema = train_op_with_ema_dependency(
self.discriminator_vars, self.train_disc)
with tf.variable_scope("Latent_variable"):
self.z_optim = tf.get_variable(
name="z_optim",
shape=[
self.config.data_loader.test_batch,
self.config.trainer.noise_dim
],
initializer=tf.truncated_normal_initializer(),
)
reinit_z = self.z_optim.initializer
with tf.variable_scope("ANOGAN"):
with tf.variable_scope("Generator_Model"):
self.x_gen_ema = self.generator(self.noise_tensor,
getter=sn.get_getter(
self.gen_ema))
self.rec_gen_ema = self.generator(self.z_optim,
getter=sn.get_getter(
self.gen_ema))
# Pass real and fake images into discriminator separately
with tf.variable_scope("Discriminator_Model"):
real_d_ema, inter_layer_real_ema = self.discriminator(
self.image_input, getter=sn.get_getter(self.dis_ema))
fake_d_ema, inter_layer_fake_ema = self.discriminator(
self.rec_gen_ema, getter=sn.get_getter(self.dis_ema))
with tf.variable_scope("Testing"):
with tf.variable_scope("Reconstruction_Loss"):
delta = self.image_input - self.rec_gen_ema
delta_flat = tf.layers.Flatten()(delta)
self.reconstruction_score_1 = tf.norm(
delta_flat,
ord=1,
axis=1,
keepdims=False,
name="epsilon",
)
self.reconstruction_score_2 = tf.norm(
delta_flat,
ord=2,
axis=1,
keepdims=False,
name="epsilon",
)
with tf.variable_scope("Discriminator_Scores"):
if self.config.trainer.loss_method == "c_entropy":
dis_score = tf.nn.sigmoid_cross_entropy_with_logits(
labels=tf.ones_like(fake_d_ema), logits=fake_d_ema)
elif self.config.trainer.loss_method == "fm":
fm = inter_layer_real_ema - inter_layer_fake_ema
fm = tf.layers.Flatten()(fm)
dis_score = tf.norm(fm,
ord=self.config.trainer.degree,
axis=1,
keepdims=False,
name="d_loss")
self.dis_score = tf.squeeze(dis_score)
with tf.variable_scope("Score"):
self.loss_invert_1 = (
self.config.trainer.weight * self.reconstruction_score_1 +
(1 - self.config.trainer.weight) * self.dis_score)
self.loss_invert_2 = (
self.config.trainer.weight * self.reconstruction_score_2 +
(1 - self.config.trainer.weight) * self.dis_score)
self.rec_error_valid = tf.reduce_mean(self.gen_loss)
with tf.variable_scope("Test_Learning_Rate"):
step_lr = tf.Variable(0, trainable=False)
learning_rate_invert = 0.001
reinit_lr = tf.variables_initializer(
tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES,
scope="Test_Learning_Rate"))
with tf.name_scope("Test_Optimizer"):
self.invert_op = tf.train.AdamOptimizer(
learning_rate_invert).minimize(self.loss_invert_1,
global_step=step_lr,
var_list=[self.z_optim],
name="optimizer")
reinit_optim = tf.variables_initializer(
tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES,
scope="Test_Optimizer"))
self.reinit_test_graph_op = [reinit_z, reinit_lr, reinit_optim]
with tf.name_scope("Scores"):
self.list_scores = self.loss_invert_1
if self.config.log.enable_summary:
with tf.name_scope("Training_Summary"):
with tf.name_scope("Dis_Summary"):
tf.summary.scalar("Real_Discriminator_Loss",
self.disc_loss_real, ["dis"])
tf.summary.scalar("Fake_Discriminator_Loss",
self.disc_loss_fake, ["dis"])
tf.summary.scalar("Discriminator_Loss",
self.total_disc_loss, ["dis"])
with tf.name_scope("Gen_Summary"):
tf.summary.scalar("Loss_Generator", self.gen_loss, ["gen"])
with tf.name_scope("Img_Summary"):
heatmap_pl_latent = tf.placeholder(tf.float32,
shape=(1, 480, 640, 3),
name="heatmap_pl_latent")
sum_op_latent = tf.summary.image("heatmap_latent",
heatmap_pl_latent)
with tf.name_scope("Validation_Summary"):
tf.summary.scalar("valid", self.rec_error_valid, ["v"])
with tf.name_scope("image_summary"):
tf.summary.image("reconstruct", self.img_gen, 1, ["image"])
tf.summary.image("input_images", self.image_input, 1,
["image"])
tf.summary.image("reconstruct", self.rec_gen_ema, 1,
["image_2"])
tf.summary.image("input_images", self.image_input, 1,
["image_2"])
self.sum_op_dis = tf.summary.merge_all("dis")
self.sum_op_gen = tf.summary.merge_all("gen")
self.sum_op = tf.summary.merge([self.sum_op_dis, self.sum_op_gen])
self.sum_op_im = tf.summary.merge_all("image")
self.sum_op_im_test = tf.summary.merge_all("image_2")
self.sum_op_valid = tf.summary.merge_all("v")
def build_model(self):
############################################################################################
# INIT
############################################################################################
# Kernel initialization for the convolutions
if self.config.trainer.init_type == "normal":
self.init_kernel = tf.random_normal_initializer(mean=0.0, stddev=0.02)
elif self.config.trainer.init_type == "xavier":
self.init_kernel = tf.contrib.layers.xavier_initializer(
uniform=False, seed=None, dtype=tf.float32
)
# Placeholders
self.is_training_gen = tf.placeholder(tf.bool)
self.is_training_dis = tf.placeholder(tf.bool)
self.is_training_enc_g = tf.placeholder(tf.bool)
self.is_training_enc_r = tf.placeholder(tf.bool)
self.image_input = tf.placeholder(
tf.float32, shape=[None] + self.config.trainer.image_dims, name="x"
)
self.noise_tensor = tf.placeholder(
tf.float32, shape=[None, self.config.trainer.noise_dim], name="noise"
)
self.denoiser_noise = tf.placeholder(
tf.float32, shape=[None] + self.config.trainer.image_dims, name="noise"
)
############################################################################################
# MODEL
############################################################################################
self.logger.info("Building training graph...")
with tf.variable_scope("SENCEBGAN_Denoiser"):
# First training part
# G(z) ==> x'
with tf.variable_scope("Generator_Model"):
self.image_gen = self.generator(self.noise_tensor)
# Discriminator outputs
with tf.variable_scope("Discriminator_Model"):
self.embedding_real, self.decoded_real = self.discriminator(
self.image_input, do_spectral_norm=self.config.trainer.do_spectral_norm
)
self.embedding_fake, self.decoded_fake = self.discriminator(
self.image_gen, do_spectral_norm=self.config.trainer.do_spectral_norm
)
# Second training part
# E(x) ==> z'
with tf.variable_scope("Encoder_G_Model"):
self.image_encoded = self.encoder_g(self.image_input)
# G(z') ==> G(E(x)) ==> x''
with tf.variable_scope("Generator_Model"):
self.image_gen_enc = self.generator(self.image_encoded)
# Discriminator outputs
with tf.variable_scope("Discriminator_Model"):
self.embedding_enc_fake, self.decoded_enc_fake = self.discriminator(
self.image_gen_enc, do_spectral_norm=self.config.trainer.do_spectral_norm
)
self.embedding_enc_real, self.decoded_enc_real = self.discriminator(
self.image_input, do_spectral_norm=self.config.trainer.do_spectral_norm
)
with tf.variable_scope("Discriminator_Model_XX"):
self.im_logit_real, self.im_f_real = self.discriminator_xx(
self.image_input,
self.image_input,
do_spectral_norm=self.config.trainer.do_spectral_norm,
)
self.im_logit_fake, self.im_f_fake = self.discriminator_xx(
self.image_input,
self.image_gen_enc,
do_spectral_norm=self.config.trainer.do_spectral_norm,
)
# Third training part
with tf.variable_scope("Encoder_G_Model"):
self.image_encoded_r = self.encoder_g(self.image_input)
with tf.variable_scope("Generator_Model"):
self.image_gen_enc_r = self.generator(self.image_encoded_r)
with tf.variable_scope("Denoiser_Model"):
self.mask, self.output = self.denoiser(self.image_gen_enc_r + self.denoiser_noise)
############################################################################################
# LOSS FUNCTIONS
############################################################################################
with tf.name_scope("Loss_Functions"):
with tf.name_scope("Generator_Discriminator"):
# Discriminator Loss
if self.config.trainer.mse_mode == "norm":
self.disc_loss_real = tf.reduce_mean(
self.mse_loss(
self.decoded_real,
self.image_input,
mode="norm",
order=self.config.trainer.order,
)
)
self.disc_loss_fake = tf.reduce_mean(
self.mse_loss(
self.decoded_fake,
self.image_gen,
mode="norm",
order=self.config.trainer.order,
)
)
elif self.config.trainer.mse_mode == "mse":
self.disc_loss_real = self.mse_loss(
self.decoded_real,
self.image_input,
mode="mse",
order=self.config.trainer.order,
)
self.disc_loss_fake = self.mse_loss(
self.decoded_fake,
self.image_gen,
mode="mse",
order=self.config.trainer.order,
)
self.loss_discriminator = (
tf.math.maximum(self.config.trainer.disc_margin - self.disc_loss_fake, 0)
+ self.disc_loss_real
)
# Generator Loss
pt_loss = 0
if self.config.trainer.pullaway:
pt_loss = self.pullaway_loss(self.embedding_fake)
self.loss_generator = self.disc_loss_fake + self.config.trainer.pt_weight * pt_loss
# New addition to enforce visual similarity
delta_noise = self.embedding_real - self.embedding_fake
delta_flat = tf.layers.Flatten()(delta_noise)
loss_noise_gen = tf.reduce_mean(
tf.norm(delta_flat, ord=2, axis=1, keepdims=False)
)
self.loss_generator += (0.1 * loss_noise_gen)
with tf.name_scope("Encoder_G"):
if self.config.trainer.mse_mode == "norm":
self.loss_enc_rec = tf.reduce_mean(
self.mse_loss(
self.image_gen_enc,
self.image_input,
mode="norm",
order=self.config.trainer.order,
)
)
self.loss_enc_f = tf.reduce_mean(
self.mse_loss(
self.embedding_enc_real,
self.embedding_enc_fake,
mode="norm",
order=self.config.trainer.order,
)
)
elif self.config.trainer.mse_mode == "mse":
self.loss_enc_rec = tf.reduce_mean(
self.mse_loss(
self.image_gen_enc,
self.image_input,
mode="mse",
order=self.config.trainer.order,
)
)
self.loss_enc_f = tf.reduce_mean(
self.mse_loss(
self.embedding_enc_real,
self.embedding_enc_fake,
mode="mse",
order=self.config.trainer.order,
)
)
self.loss_encoder_g = (
self.loss_enc_rec + self.config.trainer.encoder_f_factor * self.loss_enc_f
)
if self.config.trainer.enable_disc_xx:
self.enc_xx_real = tf.nn.sigmoid_cross_entropy_with_logits(
logits=self.im_logit_real, labels=tf.zeros_like(self.im_logit_real)
)
self.enc_xx_fake = tf.nn.sigmoid_cross_entropy_with_logits(
logits=self.im_logit_fake, labels=tf.ones_like(self.im_logit_fake)
)
self.enc_loss_xx = tf.reduce_mean(self.enc_xx_real + self.enc_xx_fake)
self.loss_encoder_g += self.enc_loss_xx
with tf.name_scope("Denoiser"):
# This part was normally the rec_image which was input but trying sth new.
delta_den = self.output - self.image_input
delta_den = tf.layers.Flatten()(delta_den)
self.den_loss = tf.reduce_mean(
tf.norm(
delta_den,
ord=2,
axis=1,
keepdims=False,
)
)
if self.config.trainer.enable_disc_xx:
with tf.name_scope("Discriminator_XX"):
self.loss_xx_real = tf.nn.sigmoid_cross_entropy_with_logits(
logits=self.im_logit_real, labels=tf.ones_like(self.im_logit_real)
)
self.loss_xx_fake = tf.nn.sigmoid_cross_entropy_with_logits(
logits=self.im_logit_fake, labels=tf.zeros_like(self.im_logit_fake)
)
self.dis_loss_xx = tf.reduce_mean(self.loss_xx_real + self.loss_xx_fake)
if self.config.trainer.enable_disc_zz:
with tf.name_scope("Discriminator_ZZ"):
self.loss_zz_real = tf.nn.sigmoid_cross_entropy_with_logits(
logits=self.z_logit_real, labels=tf.ones_like(self.z_logit_real)
)
self.loss_zz_fake = tf.nn.sigmoid_cross_entropy_with_logits(
logits=self.z_logit_fake, labels=tf.zeros_like(self.z_logit_fake)
)
self.dis_loss_zz = tf.reduce_mean(self.loss_zz_real + self.loss_zz_fake)
############################################################################################
# OPTIMIZERS
############################################################################################
with tf.name_scope("Optimizers"):
self.generator_optimizer = tf.train.AdamOptimizer(
self.config.trainer.standard_lr_gen,
beta1=self.config.trainer.optimizer_adam_beta1,
beta2=self.config.trainer.optimizer_adam_beta2,
)
self.encoder_g_optimizer = tf.train.AdamOptimizer(
self.config.trainer.standard_lr_enc,
beta1=self.config.trainer.optimizer_adam_beta1,
beta2=self.config.trainer.optimizer_adam_beta2,
)
self.denoiser_optimizer = tf.train.AdamOptimizer(
1e-4,
beta1=self.config.trainer.optimizer_adam_beta1,
beta2=self.config.trainer.optimizer_adam_beta2,
)
self.discriminator_optimizer = tf.train.AdamOptimizer(
self.config.trainer.standard_lr_dis,
beta1=self.config.trainer.optimizer_adam_beta1,
beta2=self.config.trainer.optimizer_adam_beta2,
)
# Collect all the variables
all_variables = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES)
# Generator Network Variables
self.generator_vars = [
v for v in all_variables if v.name.startswith("SENCEBGAN_Denoiser/Generator_Model")
]
# Discriminator Network Variables
self.discriminator_vars = [
v for v in all_variables if v.name.startswith("SENCEBGAN_Denoiser/Discriminator_Model")
]
# Discriminator Network Variables
self.encoder_g_vars = [
v for v in all_variables if v.name.startswith("SENCEBGAN_Denoiser/Encoder_G_Model")
]
self.denoiser_vars = [
v for v in all_variables if v.name.startswith("SENCEBGAN_Denoiser/Denoiser_Model")
]
self.dxxvars = [
v for v in all_variables if v.name.startswith("SENCEBGAN_Denoiser/Discriminator_Model_XX")
]
self.dzzvars = [
v for v in all_variables if v.name.startswith("SENCEBGAN_Denoiser/Discriminator_Model_ZZ")
]
# Generator Network Operations
self.gen_update_ops = tf.get_collection(
tf.GraphKeys.UPDATE_OPS, scope="SENCEBGAN_Denoiser/Generator_Model"
)
# Discriminator Network Operations
self.disc_update_ops = tf.get_collection(
tf.GraphKeys.UPDATE_OPS, scope="SENCEBGAN_Denoiser/Discriminator_Model"
)
self.encg_update_ops = tf.get_collection(
tf.GraphKeys.UPDATE_OPS, scope="SENCEBGAN_Denoiser/Encoder_G_Model"
)
self.den_update_ops = tf.get_collection(
tf.GraphKeys.UPDATE_OPS, scope="SENCEBGAN_Denoiser/Denoiser_Model"
)
self.update_ops_dis_xx = tf.get_collection(
tf.GraphKeys.UPDATE_OPS, scope="SENCEBGAN_Denoiser/Discriminator_Model_XX"
)
self.update_ops_dis_zz = tf.get_collection(
tf.GraphKeys.UPDATE_OPS, scope="SENCEBGAN_Denoiser/Discriminator_Model_ZZ"
)
with tf.control_dependencies(self.gen_update_ops):
self.gen_op = self.generator_optimizer.minimize(
self.loss_generator,
var_list=self.generator_vars,
global_step=self.global_step_tensor,
)
with tf.control_dependencies(self.disc_update_ops):
self.disc_op = self.discriminator_optimizer.minimize(
self.loss_discriminator, var_list=self.discriminator_vars
)
with tf.control_dependencies(self.encg_update_ops):
self.encg_op = self.encoder_g_optimizer.minimize(
self.loss_encoder_g,
var_list=self.encoder_g_vars,
global_step=self.global_step_tensor,
)
with tf.control_dependencies(self.den_update_ops):
self.den_op = self.denoiser_optimizer.minimize(
self.den_loss,
var_list=self.denoiser_vars,
global_step=self.global_step_tensor,
)
if self.config.trainer.enable_disc_xx:
with tf.control_dependencies(self.update_ops_dis_xx):
self.disc_op_xx = self.discriminator_optimizer.minimize(
self.dis_loss_xx, var_list=self.dxxvars
)
if self.config.trainer.enable_disc_zz:
with tf.control_dependencies(self.update_ops_dis_zz):
self.disc_op_zz = self.discriminator_optimizer.minimize(
self.dis_loss_zz, var_list=self.dzzvars
)
# if self.config.trainer.extra_gan_training:
# with tf.control_dependencies(self.gen_update_ops):
# self.gen_2_op = self.generator_optimizer.minimize(
# self.loss_generator_2, var_list=self.generator_vars
# )
# Exponential Moving Average for Estimation
self.dis_ema = tf.train.ExponentialMovingAverage(decay=self.config.trainer.ema_decay)
maintain_averages_op_dis = self.dis_ema.apply(self.discriminator_vars)
self.gen_ema = tf.train.ExponentialMovingAverage(decay=self.config.trainer.ema_decay)
maintain_averages_op_gen = self.gen_ema.apply(self.generator_vars)
# if self.config.trainer.extra_gan_training:
# self.gen_2_ema = tf.train.ExponentialMovingAverage(decay=self.config.trainer.ema_decay)
# maintain_averages_op_gen_2 = self.gen_2_ema.apply(self.generator_vars)
self.encg_ema = tf.train.ExponentialMovingAverage(decay=self.config.trainer.ema_decay)
maintain_averages_op_encg = self.encg_ema.apply(self.encoder_g_vars)
self.den_ema = tf.train.ExponentialMovingAverage(decay=self.config.trainer.ema_decay)
maintain_averages_op_den = self.den_ema.apply(self.denoiser_vars)
if self.config.trainer.enable_disc_xx:
self.dis_xx_ema = tf.train.ExponentialMovingAverage(
decay=self.config.trainer.ema_decay
)
maintain_averages_op_dis_xx = self.dis_xx_ema.apply(self.dxxvars)
if self.config.trainer.enable_disc_zz:
self.dis_zz_ema = tf.train.ExponentialMovingAverage(
decay=self.config.trainer.ema_decay
)
maintain_averages_op_dis_zz = self.dis_zz_ema.apply(self.dzzvars)
with tf.control_dependencies([self.disc_op]):
self.train_dis_op = tf.group(maintain_averages_op_dis)
with tf.control_dependencies([self.gen_op]):
self.train_gen_op = tf.group(maintain_averages_op_gen)
# if self.config.trainer.extra_gan_training:
# with tf.control_dependencies([self.gen_2_op]):
# self.train_gen_op_2 = tf.group(maintain_averages_op_gen_2)
with tf.control_dependencies([self.encg_op]):
self.train_enc_g_op = tf.group(maintain_averages_op_encg)
with tf.control_dependencies([self.den_op]):
self.train_den_op = tf.group(maintain_averages_op_den)
if self.config.trainer.enable_disc_xx:
with tf.control_dependencies([self.disc_op_xx]):
self.train_dis_op_xx = tf.group(maintain_averages_op_dis_xx)
if self.config.trainer.enable_disc_zz:
with tf.control_dependencies([self.disc_op_zz]):
self.train_dis_op_zz = tf.group(maintain_averages_op_dis_zz)
############################################################################################
# TESTING
############################################################################################
self.logger.info("Building Testing Graph...")
with tf.variable_scope("SENCEBGAN_Denoiser"):
with tf.variable_scope("Discriminator_Model"):
self.embedding_q_ema, self.decoded_q_ema = self.discriminator(
self.image_input,
getter=get_getter(self.dis_ema),
do_spectral_norm=self.config.trainer.do_spectral_norm,
)
with tf.variable_scope("Generator_Model"):
self.image_gen_ema = self.generator(
self.embedding_q_ema, getter=get_getter(self.gen_ema)
)
with tf.variable_scope("Discriminator_Model"):
self.embedding_rec_ema, self.decoded_rec_ema = self.discriminator(
self.image_gen_ema,
getter=get_getter(self.dis_ema),
do_spectral_norm=self.config.trainer.do_spectral_norm,
)
# Second Training Part
with tf.variable_scope("Encoder_G_Model"):
self.image_encoded_ema = self.encoder_g(
self.image_input, getter=get_getter(self.encg_ema)
)
with tf.variable_scope("Generator_Model"):
self.image_gen_enc_ema = self.generator(
self.image_encoded_ema, getter=get_getter(self.gen_ema)
)
with tf.variable_scope("Discriminator_Model"):
self.embedding_enc_fake_ema, self.decoded_enc_fake_ema = self.discriminator(
self.image_gen_enc_ema,
getter=get_getter(self.dis_ema),
do_spectral_norm=self.config.trainer.do_spectral_norm,
)
self.embedding_enc_real_ema, self.decoded_enc_real_ema = self.discriminator(
self.image_input,
getter=get_getter(self.dis_ema),
do_spectral_norm=self.config.trainer.do_spectral_norm,
)
if self.config.trainer.enable_disc_xx:
with tf.variable_scope("Discriminator_Model_XX"):
self.im_logit_real_ema, self.im_f_real_ema = self.discriminator_xx(
self.image_input,
self.image_input,
getter=get_getter(self.dis_xx_ema),
do_spectral_norm=self.config.trainer.do_spectral_norm,
)
self.im_logit_fake_ema, self.im_f_fake_ema = self.discriminator_xx(
self.image_input,
self.image_gen_enc_ema,
getter=get_getter(self.dis_xx_ema),
do_spectral_norm=self.config.trainer.do_spectral_norm,
)
# Third training part
with tf.variable_scope("Encoder_G_Model"):
self.image_encoded_r_ema = self.encoder_g(self.image_input)
with tf.variable_scope("Generator_Model"):
self.image_gen_enc_r_ema = self.generator(self.image_encoded_r_ema)
with tf.variable_scope("Denoiser_Model"):
self.output_ema, self.mask_ema = self.denoiser(self.image_gen_enc_r_ema, getter=get_getter(self.den_ema))
with tf.name_scope("Testing"):
with tf.name_scope("Image_Based"):
delta = self.image_input - self.image_gen_enc_ema
self.mask = -delta
delta_flat = tf.layers.Flatten()(delta)
img_score_l1 = tf.norm(
delta_flat, ord=1, axis=1, keepdims=False, name="img_loss__1"
)
self.img_score_l1 = tf.squeeze(img_score_l1)
delta = self.embedding_enc_fake_ema - self.embedding_enc_real_ema
delta_flat = tf.layers.Flatten()(delta)
img_score_l2 = tf.norm(
delta_flat, ord=1, axis=1, keepdims=False, name="img_loss__2"
)
self.img_score_l2 = tf.squeeze(img_score_l2)
self.score_comb = (
(1 - self.config.trainer.feature_match_weight) * self.img_score_l1
+ self.config.trainer.feature_match_weight * self.img_score_l2
)
with tf.variable_scope("Pipeline_Loss_1"):
delta_pipe = self.output_ema - self.image_input
delta_pipe = tf.layers.Flatten()(delta_pipe)
self.pipe_score = tf.norm(delta_pipe, ord=1,axis=1,keepdims=False)
with tf.variable_scope("Pipeline_Loss_2"):
delta_pipe = self.output_ema - self.image_input
delta_pipe = tf.layers.Flatten()(delta_pipe)
self.pipe_score_2 = tf.norm(delta_pipe, ord=2,axis=1,keepdims=False)
with tf.name_scope("Noise_Based"):
with tf.variable_scope("Mask_1"):
delta_mask = (self.image_input - self.mask_ema)
delta_mask = tf.layers.Flatten()(delta_mask)
self.mask_score_1 = tf.norm(delta_mask, ord=1,axis=1,keepdims=False)
with tf.variable_scope("Mask_2"):
delta_mask_2 = (self.image_input - self.mask_ema)
delta_mask_2 = tf.layers.Flatten()(delta_mask_2)
self.mask_score_2 = tf.norm(delta_mask_2, ord=2,axis=1,keepdims=False)
############################################################################################
# TENSORBOARD
############################################################################################
if self.config.log.enable_summary:
with tf.name_scope("train_summary"):
with tf.name_scope("dis_summary"):
tf.summary.scalar("loss_disc", self.loss_discriminator, ["dis"])
tf.summary.scalar("loss_disc_real", self.disc_loss_real, ["dis"])
tf.summary.scalar("loss_disc_fake", self.disc_loss_fake, ["dis"])
if self.config.trainer.enable_disc_xx:
tf.summary.scalar("loss_dis_xx", self.dis_loss_xx, ["enc_g"])
if self.config.trainer.enable_disc_zz:
tf.summary.scalar("loss_dis_zz", self.dis_loss_zz, ["den"])
with tf.name_scope("gen_summary"):
tf.summary.scalar("loss_generator", self.loss_generator, ["gen"])
with tf.name_scope("enc_summary"):
tf.summary.scalar("loss_encoder_g", self.loss_encoder_g, ["enc_g"])
tf.summary.scalar("loss_den", self.den_loss, ["den"])
with tf.name_scope("img_summary"):
tf.summary.image("input_image", self.image_input, 1, ["img_1"])
tf.summary.image("reconstructed", self.image_gen, 1, ["img_1"])
tf.summary.image("input_enc", self.image_input, 1, ["img_2"])
tf.summary.image("reconstructed", self.image_gen_enc, 1, ["img_2"])
tf.summary.image("input_image",self.image_input,1,["test"])
tf.summary.image("reconstructed", self.image_gen_enc_r_ema,1,["test"])
tf.summary.image("mask", self.mask_ema, 1, ["test"])
tf.summary.image("output", self.output_ema, 1, ["test"])
self.sum_op_dis = tf.summary.merge_all("dis")
self.sum_op_gen = tf.summary.merge_all("gen")
self.sum_op_enc_g = tf.summary.merge_all("enc_g")
self.sum_op_den = tf.summary.merge_all("den")
self.sum_op_im_1 = tf.summary.merge_all("img_1")
self.sum_op_im_2 = tf.summary.merge_all("img_2")
self.sum_op_test = tf.summary.merge_all("test")
self.sum_op = tf.summary.merge([self.sum_op_dis, self.sum_op_gen])
def build_model(self):
# Place holders
self.img_size = self.config.data_loader.image_size
self.is_training = tf.placeholder(tf.bool)
self.image_input = tf.placeholder(dtype=tf.float32,
shape=[None] +
self.config.trainer.image_dims,
name="x")
if self.config.trainer.init_type == "normal":
self.init_kernel = tf.random_normal_initializer(mean=0.0,
stddev=0.02)
elif self.config.trainer.init_type == "xavier":
self.init_kernel = tf.contrib.layers.xavier_initializer(
uniform=False, seed=None, dtype=tf.float32)
self.true_labels = tf.placeholder(dtype=tf.float32,
shape=[None, 1],
name="true_labels")
self.generated_labels = tf.placeholder(dtype=tf.float32,
shape=[None, 1],
name="gen_labels")
self.real_noise = tf.placeholder(dtype=tf.float32,
shape=[None] +
self.config.trainer.image_dims,
name="real_noise")
self.fake_noise = tf.placeholder(dtype=tf.float32,
shape=[None] +
self.config.trainer.image_dims,
name="fake_noise")
#######################################################################
# GRAPH
########################################################################
self.logger.info("Building Training Graph")
with tf.variable_scope("Skip_GANomaly"):
with tf.variable_scope("Generator_Model"):
self.img_rec = self.generator(self.image_input)
self.img_rec += self.fake_noise
with tf.variable_scope("Discriminator_Model"):
self.disc_real, self.inter_layer_real = self.discriminator(
self.image_input + self.real_noise)
self.disc_fake, self.inter_layer_fake = self.discriminator(
self.img_rec)
########################################################################
# METRICS
########################################################################
with tf.variable_scope("Loss_Functions"):
with tf.variable_scope("Discriminator_Loss"):
# According to the paper we invert the values for the normal/fake.
# So normal images should be labeled
# as zeros
self.loss_dis_real = tf.reduce_mean(
tf.nn.sigmoid_cross_entropy_with_logits(
labels=self.true_labels, logits=self.disc_real))
self.loss_dis_fake = tf.reduce_mean(
tf.nn.sigmoid_cross_entropy_with_logits(
labels=self.generated_labels, logits=self.disc_fake))
# Adversarial Loss Part for Discriminator
self.loss_discriminator = self.loss_dis_real + self.loss_dis_fake
with tf.variable_scope("Generator_Loss"):
# Adversarial Loss
if self.config.trainer.flip_labels:
labels = tf.zeros_like(self.disc_fake)
else:
labels = tf.ones_like(self.disc_fake)
self.gen_adv_loss = tf.reduce_mean(
tf.nn.sigmoid_cross_entropy_with_logits(
logits=self.disc_fake, labels=labels))
# Contextual Loss
context_layers = self.image_input - self.img_rec
self.contextual_loss = tf.reduce_mean(
tf.norm(context_layers,
ord=1,
axis=1,
keepdims=False,
name="Contextual_Loss"))
# Latent Loss
layer_diff = self.inter_layer_real - self.inter_layer_fake
self.latent_loss = tf.reduce_mean(
tf.norm(layer_diff,
ord=2,
axis=1,
keepdims=False,
name="Latent_Loss"))
self.gen_loss_total = (
self.config.trainer.weight_adv * self.gen_adv_loss +
self.config.trainer.weight_cont * self.contextual_loss +
self.config.trainer.weight_lat * self.latent_loss)
########################################################################
# OPTIMIZATION
########################################################################
# Build the Optimizers
with tf.name_scope("Optimization"):
self.generator_optimizer = tf.train.AdamOptimizer(
self.config.trainer.generator_l_rate,
beta1=self.config.trainer.optimizer_adam_beta1,
beta2=self.config.trainer.optimizer_adam_beta2,
)
self.discriminator_optimizer = tf.train.AdamOptimizer(
self.config.trainer.discriminator_l_rate,
beta1=self.config.trainer.optimizer_adam_beta1,
beta2=self.config.trainer.optimizer_adam_beta2,
)
# Collect all the variables
all_variables = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES)
# Generator Network Variables
self.generator_vars = [
v for v in all_variables
if v.name.startswith("Skip_GANomaly/Generator_Model")
]
# Discriminator Network Variables
self.discriminator_vars = [
v for v in all_variables
if v.name.startswith("Skip_GANomaly/Discriminator_Model")
]
# Generator Network Operations
self.gen_update_ops = tf.get_collection(
tf.GraphKeys.UPDATE_OPS, scope="Skip_GANomaly/Generator_Model")
# Discriminator Network Operations
self.disc_update_ops = tf.get_collection(
tf.GraphKeys.UPDATE_OPS,
scope="Skip_GANomaly/Discriminator_Model")
# Initialization of Optimizers
with tf.control_dependencies(self.gen_update_ops):
self.gen_op = self.generator_optimizer.minimize(
self.gen_loss_total,
global_step=self.global_step_tensor,
var_list=self.generator_vars,
)
with tf.control_dependencies(self.disc_update_ops):
self.disc_op = self.discriminator_optimizer.minimize(
self.loss_discriminator, var_list=self.discriminator_vars)
# Exponential Moving Average for Estimation
self.dis_ema = tf.train.ExponentialMovingAverage(
decay=self.config.trainer.ema_decay)
maintain_averages_op_dis = self.dis_ema.apply(
self.discriminator_vars)
self.gen_ema = tf.train.ExponentialMovingAverage(
decay=self.config.trainer.ema_decay)
maintain_averages_op_gen = self.gen_ema.apply(self.generator_vars)
with tf.control_dependencies([self.disc_op]):
self.train_dis_op = tf.group(maintain_averages_op_dis)
with tf.control_dependencies([self.gen_op]):
self.train_gen_op = tf.group(maintain_averages_op_gen)
########################################################################
# TESTING
########################################################################
self.logger.info("Building Testing Graph...")
with tf.variable_scope("Skip_GANomaly"):
with tf.variable_scope("Generator_Model"):
self.img_rec_ema = self.generator(self.image_input,
getter=get_getter(
self.gen_ema))
with tf.variable_scope("Discriminator_Model"):
self.disc_real_ema, self.inter_layer_real_ema = self.discriminator(
self.image_input, getter=get_getter(self.dis_ema))
self.disc_fake_ema, self.inter_layer_fake_ema = self.discriminator(
self.img_rec_ema, getter=get_getter(self.dis_ema))
with tf.name_scope("Testing"):
with tf.variable_scope("Reconstruction_Loss"):
# Contextual Loss
context_layers = self.image_input - self.img_rec_ema
context_layers = tf.layers.Flatten()(context_layers)
contextual_loss_ema_1 = tf.norm(context_layers,
ord=1,
axis=1,
keepdims=False,
name="Contextual_Loss")
self.contextual_loss_ema_1 = tf.squeeze(contextual_loss_ema_1)
contextual_loss_ema_2 = tf.norm(context_layers,
ord=2,
axis=1,
keepdims=False,
name="Contextual_Loss")
self.contextual_loss_ema_2 = tf.squeeze(contextual_loss_ema_2)
with tf.variable_scope("Latent_Loss"):
# Latent Loss
layer_diff = self.inter_layer_real_ema - self.inter_layer_fake_ema
layer_diff = tf.layers.Flatten()(layer_diff)
latent_loss_ema = tf.norm(layer_diff,
ord=2,
axis=None,
keepdims=False,
name="Latent_Loss")
self.latent_loss_ema = tf.squeeze(latent_loss_ema)
self.anomaly_score_1 = (
self.config.trainer.weight * self.contextual_loss_ema_1 +
(1 - self.config.trainer.weight) * self.latent_loss_ema)
self.anomaly_score_2 = (
self.config.trainer.weight * self.contextual_loss_ema_2 +
(1 - self.config.trainer.weight) * self.latent_loss_ema)
if self.config.trainer.enable_early_stop:
self.rec_error_valid = tf.reduce_mean(self.latent_loss_ema)
########################################################################
# TENSORBOARD
########################################################################
if self.config.log.enable_summary:
with tf.name_scope("summary"):
with tf.name_scope("disc_summary"):
tf.summary.scalar("loss_discriminator_total",
self.loss_discriminator, ["dis"])
tf.summary.scalar("loss_dis_real", self.loss_dis_real,
["dis"])
tf.summary.scalar("loss_dis_fake", self.loss_dis_fake,
["dis"])
with tf.name_scope("gen_summary"):
tf.summary.scalar("loss_generator_total",
self.gen_loss_total, ["gen"])
tf.summary.scalar("loss_gen_adv", self.gen_adv_loss,
["gen"])
tf.summary.scalar("loss_gen_con", self.contextual_loss,
["gen"])
tf.summary.scalar("loss_gen_enc", self.latent_loss,
["gen"])
with tf.name_scope("image_summary"):
tf.summary.image("reconstruct", self.img_rec, 1, ["image"])
tf.summary.image("input_images", self.image_input, 1,
["image"])
tf.summary.image("reconstruct", self.img_rec_ema, 1,
["image_2"])
tf.summary.image("input_images", self.image_input, 1,
["image_2"])
if self.config.trainer.enable_early_stop:
with tf.name_scope("validation_summary"):
tf.summary.scalar("valid", self.rec_error_valid, ["v"])
self.sum_op_dis = tf.summary.merge_all("dis")
self.sum_op_gen = tf.summary.merge_all("gen")
self.sum_op_im = tf.summary.merge_all("image")
self.sum_op_im_test = tf.summary.merge_all("image_2")
self.sum_op_valid = tf.summary.merge_all("v")
def build_model(self):
# Kernel initialization for the convolutions
self.init_kernel = tf.random_normal_initializer(mean=0.0, stddev=0.02)
# Placeholders
self.is_training = tf.placeholder(tf.bool)
self.image_input = tf.placeholder(tf.float32,
shape=[None] +
self.config.trainer.image_dims,
name="x")
self.noise_tensor = tf.placeholder(
tf.float32,
shape=[None, self.config.trainer.noise_dim],
name="noise")
self.true_labels = tf.placeholder(dtype=tf.float32,
shape=[None, 1],
name="true_labels")
self.generated_labels = tf.placeholder(dtype=tf.float32,
shape=[None, 1],
name="gen_labels")
self.real_noise = tf.placeholder(dtype=tf.float32,
shape=[None] +
self.config.trainer.image_dims,
name="real_noise")
self.fake_noise = tf.placeholder(dtype=tf.float32,
shape=[None] +
self.config.trainer.image_dims,
name="fake_noise")
self.logger.info("Building training graph...")
with tf.variable_scope("BIGAN"):
with tf.variable_scope("Encoder_Model"):
self.noise_gen = self.encoder(
self.image_input,
do_spectral_norm=self.config.trainer.do_spectral_norm)
with tf.variable_scope("Generator_Model"):
self.image_gen = self.generator(
self.noise_tensor) + self.fake_noise
self.reconstructed = self.generator(self.noise_gen)
with tf.variable_scope("Discriminator_Model"):
# E(x) and x --> This being real is the output of discriminator
l_encoder, inter_layer_inp = self.discriminator(
self.noise_gen,
self.image_input + self.real_noise,
do_spectral_norm=self.config.trainer.do_spectral_norm,
)
# z and G(z)
l_generator, inter_layer_rct = self.discriminator(
self.noise_tensor,
self.image_gen,
do_spectral_norm=self.config.trainer.do_spectral_norm,
)
# Loss Function Implementations
with tf.name_scope("Loss_Functions"):
# Discriminator
# Discriminator sees the encoder result as true because it discriminates E(x), x as the real pair
if self.config.trainer.mode == "standard":
self.loss_dis_enc = tf.reduce_mean(
tf.nn.sigmoid_cross_entropy_with_logits(
labels=self.true_labels, logits=l_encoder))
self.loss_dis_gen = tf.reduce_mean(
tf.nn.sigmoid_cross_entropy_with_logits(
labels=self.generated_labels, logits=l_generator))
self.loss_discriminator = self.loss_dis_enc + self.loss_dis_gen
# Flip the weigths for the encoder and generator
if self.config.trainer.flip_labels:
labels_gen = tf.zeros_like(l_generator)
labels_enc = tf.ones_like(l_encoder)
else:
labels_gen = tf.ones_like(l_generator)
labels_enc = tf.zeros_like(l_encoder)
# Generator
# Generator is considered as the true ones here because it tries to fool discriminator
self.loss_generator_ce = tf.reduce_mean(
tf.nn.sigmoid_cross_entropy_with_logits(
labels=labels_gen, logits=l_generator))
delta = inter_layer_inp - inter_layer_rct
delta = tf.layers.Flatten()(delta)
self.loss_generator_fm = tf.reduce_mean(
tf.norm(delta, ord=2, axis=1, keepdims=False))
self.loss_generator = (
self.loss_generator_ce +
self.config.trainer.feature_match_weight *
self.loss_generator_fm)
# Encoder
# Encoder is considered as the fake one because it tries to fool the discriminator also
self.loss_encoder = tf.reduce_mean(
tf.nn.sigmoid_cross_entropy_with_logits(labels=labels_enc,
logits=l_encoder))
elif self.config.trainer.mode == "wgan":
self.loss_generator = -tf.reduce_mean(l_generator)
self.loss_encoder = -tf.reduce_mean(l_encoder)
self.loss_discriminator = tf.reduce_mean(
l_generator) - tf.reduce_mean(l_encoder)
elif self.config.trainer.mode == "wgan-gp":
self.loss_generator = -tf.reduce_mean(l_generator)
self.loss_encoder = -tf.reduce_mean(l_encoder)
self.loss_discriminator = tf.reduce_mean(
l_generator) - tf.reduce_mean(l_encoder)
alpha_x = tf.random_uniform(
shape=[self.config.data_loader.batch_size] +
self.config.trainer.image_dims,
minval=0.0,
maxval=1.0,
)
alpha_z = tf.random_uniform(
shape=[
self.config.data_loader.batch_size,
self.config.trainer.noise_dim
],
minval=0.0,
maxval=1.0,
)
differences_x = self.image_gen - self.image_input
interpolates_x = self.image_input + (alpha_x * differences_x)
differences_z = self.noise_gen - self.noise_tensor
interpolates_z = self.noise_tensor + (alpha_z * differences_z)
gradients = tf.gradients(
self.discriminator(interpolates_z, interpolates_x),
[interpolates_z, interpolates_x],
)[0]
slopes = tf.sqrt(
tf.reduce_sum(tf.square(gradients), reduction_indices=[1]))
gradient_penalty = tf.reduce_mean((slopes - 1.0)**2)
self.loss_discriminator += self.config.trainer.wgan_gp_lambda * gradient_penalty
# Optimizer Implementations
with tf.name_scope("Optimizers"):
if self.config.trainer.mode == "standard":
# Build the optimizers
self.generator_optimizer = tf.train.AdamOptimizer(
self.config.trainer.standard_lr,
beta1=self.config.trainer.optimizer_adam_beta1,
beta2=self.config.trainer.optimizer_adam_beta2,
)
self.discriminator_optimizer = tf.train.AdamOptimizer(
self.config.trainer.standard_lr,
beta1=self.config.trainer.optimizer_adam_beta1,
beta2=self.config.trainer.optimizer_adam_beta2,
)
self.encoder_optimizer = tf.train.AdamOptimizer(
self.config.trainer.standard_lr,
beta1=self.config.trainer.optimizer_adam_beta1,
beta2=self.config.trainer.optimizer_adam_beta2,
)
elif self.config.trainer.mode == "wgan":
# Build the optimizers
self.generator_optimizer = tf.train.RMSPropOptimizer(
self.config.trainer.wgan_lr)
self.discriminator_optimizer = tf.train.RMSPropOptimizer(
self.config.trainer.wgan_lr)
self.encoder_optimizer = tf.train.RMSPropOptimizer(
self.config.trainer.wgan_lr)
elif self.config.trainer.mode == "wgan-gp":
# Build the optimizers
self.generator_optimizer = tf.train.AdamOptimizer(
self.config.trainer.wgan_gp_lr,
beta1=self.config.trainer.optimizer_adam_beta1,
beta2=self.config.trainer.optimizer_adam_beta2,
)
self.discriminator_optimizer = tf.train.AdamOptimizer(
self.config.trainer.wgan_gp_lr,
beta1=self.config.trainer.optimizer_adam_beta1,
beta2=self.config.trainer.optimizer_adam_beta2,
)
self.encoder_optimizer = tf.train.AdamOptimizer(
self.config.trainer.wgan_gp_lr,
beta1=self.config.trainer.optimizer_adam_beta1,
beta2=self.config.trainer.optimizer_adam_beta2,
)
# Collect all the variables
all_variables = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES)
# Generator Network Variables
self.generator_vars = [
v for v in all_variables
if v.name.startswith("BIGAN/Generator_Model")
]
# Discriminator Network Variables
self.discriminator_vars = [
v for v in all_variables
if v.name.startswith("BIGAN/Discriminator_Model")
]
if self.config.trainer.mode == "wgan":
clip_ops = []
for var in self.discriminator_vars:
clip_bounds = [-0.01, 0.01]
clip_ops.append(
tf.assign(
var,
tf.clip_by_value(var, clip_bounds[0],
clip_bounds[1])))
self.clip_disc_weights = tf.group(*clip_ops)
# Encoder Network Variables
self.encoder_vars = [
v for v in all_variables
if v.name.startswith("BIGAN/Encoder_Model")
]
# Create Training Operations
# Generator Network Operations
self.gen_update_ops = tf.get_collection(
tf.GraphKeys.UPDATE_OPS, scope="BIGAN/Generator_Model")
# Discriminator Network Operations
self.disc_update_ops = tf.get_collection(
tf.GraphKeys.UPDATE_OPS, scope="BIGAN/Discriminator_Model")
# Encoder Network Operations
self.enc_update_ops = tf.get_collection(
tf.GraphKeys.UPDATE_OPS, scope="BIGAN/Encoder_Model")
# Initialization of Optimizers
with tf.control_dependencies(self.gen_update_ops):
self.gen_op = self.generator_optimizer.minimize(
self.loss_generator,
var_list=self.generator_vars,
global_step=self.global_step_tensor,
)
with tf.control_dependencies(self.disc_update_ops):
self.disc_op = self.discriminator_optimizer.minimize(
self.loss_discriminator, var_list=self.discriminator_vars)
with tf.control_dependencies(self.enc_update_ops):
self.enc_op = self.encoder_optimizer.minimize(
self.loss_encoder, var_list=self.encoder_vars)
# Exponential Moving Average for Estimation
self.dis_ema = tf.train.ExponentialMovingAverage(
decay=self.config.trainer.ema_decay)
maintain_averages_op_dis = self.dis_ema.apply(
self.discriminator_vars)
self.gen_ema = tf.train.ExponentialMovingAverage(
decay=self.config.trainer.ema_decay)
maintain_averages_op_gen = self.gen_ema.apply(self.generator_vars)
self.enc_ema = tf.train.ExponentialMovingAverage(
decay=self.config.trainer.ema_decay)
maintain_averages_op_enc = self.enc_ema.apply(self.encoder_vars)
with tf.control_dependencies([self.disc_op]):
self.train_dis_op = tf.group(maintain_averages_op_dis)
with tf.control_dependencies([self.gen_op]):
self.train_gen_op = tf.group(maintain_averages_op_gen)
with tf.control_dependencies([self.enc_op]):
self.train_enc_op = tf.group(maintain_averages_op_enc)
self.logger.info("Building Testing Graph...")
with tf.variable_scope("BIGAN"):
with tf.variable_scope("Encoder_Model"):
self.noise_gen_ema = self.encoder(
self.image_input,
getter=get_getter(self.enc_ema),
do_spectral_norm=self.config.trainer.do_spectral_norm,
)
with tf.variable_scope("Generator_Model"):
self.reconstruct_ema = self.generator(self.noise_gen_ema,
getter=get_getter(
self.gen_ema))
with tf.variable_scope("Discriminator_Model"):
self.l_encoder_ema, self.inter_layer_inp_ema = self.discriminator(
self.noise_gen_ema, # E(x)
self.image_input, # x
getter=get_getter(self.dis_ema),
do_spectral_norm=self.config.trainer.do_spectral_norm,
)
self.l_generator_ema, self.inter_layer_rct_ema = self.discriminator(
self.noise_gen_ema, # E(x)
self.reconstruct_ema, # G(E(x))
getter=get_getter(self.dis_ema),
do_spectral_norm=self.config.trainer.do_spectral_norm,
)
with tf.name_scope("Testing"):
with tf.variable_scope("Reconstruction_Loss"):
# LG(x) = ||x - G(E(x))||_1
delta = self.image_input - self.reconstruct_ema
delta_flat = tf.layers.Flatten()(delta)
self.gen_score = tf.norm(
delta_flat,
ord=self.config.trainer.degree,
axis=1,
keepdims=False,
name="epsilon",
)
with tf.variable_scope("Discriminator_Loss"):
if self.config.trainer.loss_method == "cross_e":
self.dis_score = tf.nn.sigmoid_cross_entropy_with_logits(
labels=tf.ones_like(self.l_encoder_ema),
logits=self.l_encoder_ema)
elif self.config.trainer.loss_method == "fm":
fm = self.inter_layer_inp_ema - self.inter_layer_rct_ema
fm = tf.layers.Flatten()(fm)
self.dis_score = tf.norm(fm,
ord=self.config.trainer.degree,
axis=1,
keepdims=False,
name="d_loss")
self.dis_score = tf.squeeze(self.dis_score)
with tf.variable_scope("Score"):
self.list_scores = (
1 - self.config.trainer.weight
) * self.dis_score + self.config.trainer.weight * self.gen_score
if self.config.trainer.enable_early_stop:
self.rec_error_valid = tf.reduce_mean(self.list_scores)
if self.config.log.enable_summary:
with tf.name_scope("Summary"):
with tf.name_scope("Disc_Summary"):
tf.summary.scalar("loss_discriminator",
self.loss_discriminator, ["dis"])
if self.config.trainer.mode == "standard":
tf.summary.scalar("loss_dis_encoder",
self.loss_dis_enc, ["dis"])
tf.summary.scalar("loss_dis_gen", self.loss_dis_gen,
["dis"])
with tf.name_scope("Gen_Summary"):
tf.summary.scalar("loss_generator", self.loss_generator,
["gen"])
if self.config.trainer.mode == "standard":
tf.summary.scalar("loss_generator_ce",
self.loss_generator_ce, ["gen"])
tf.summary.scalar("loss_generator_fm",
self.loss_generator_fm, ["gen"])
tf.summary.scalar("loss_encoder", self.loss_encoder,
["gen"])
with tf.name_scope("Image_Summary"):
tf.summary.image("reconstruct", self.reconstructed, 3,
["image"])
tf.summary.image("input_images", self.image_input, 3,
["image"])
if self.config.trainer.enable_early_stop:
with tf.name_scope("validation_summary"):
tf.summary.scalar("valid", self.rec_error_valid, ["v"])
self.sum_op_dis = tf.summary.merge_all("dis")
self.sum_op_gen = tf.summary.merge_all("gen")
self.sum_op_im = tf.summary.merge_all("image")
self.sum_op_valid = tf.summary.merge_all("v")
def build_model(self):
# Initializations
# Kernel initialization for the convolutions
self.init_kernel = tf.contrib.layers.xavier_initializer(
uniform=False, seed=None, dtype=tf.float32)
# Placeholders
self.is_training = tf.placeholder(tf.bool)
self.image_input = tf.placeholder(tf.float32,
shape=[None] +
self.config.trainer.image_dims,
name="x")
self.noise_tensor = tf.placeholder(
tf.float32,
shape=[None, self.config.trainer.noise_dim],
name="noise")
# Build Training Graph
self.logger.info("Building training graph...")
with tf.variable_scope("FenceGAN"):
with tf.variable_scope("Generator_Model"):
self.image_gen = self.generator(self.noise_tensor)
with tf.variable_scope("Discriminator_Model"):
self.disc_real, self.f_layer_real = self.discriminator(
self.image_input,
do_spectral_norm=self.config.trainer.do_spectral_norm)
self.disc_fake, self.f_layer_fake = self.discriminator(
self.image_gen,
do_spectral_norm=self.config.trainer.do_spectral_norm)
# Loss functions
with tf.name_scope("Loss_Functions"):
# Discriminator Loss
# Generator Loss
pass
# Optimizers
with tf.name_scope("Optimizers"):
self.generator_optimizer = tf.train.AdamOptimizer(
self.config.trainer.standard_lr,
beta1=self.config.trainer.optimizer_adam_beta1,
beta2=self.config.trainer.optimizer_adam_beta2,
)
self.discriminator_optimizer = tf.train.AdamOptimizer(
self.config.trainer.standard_lr,
beta1=self.config.trainer.optimizer_adam_beta1,
beta2=self.config.trainer.optimizer_adam_beta2,
)
# Collect all the variables
all_variables = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES)
# Generator Network Variables
self.generator_vars = [
v for v in all_variables
if v.name.startswith("FenceGAN/Generator_Model")
]
# Discriminator Network Variables
self.discriminator_vars = [
v for v in all_variables
if v.name.startswith("FenceGAN/Discriminator_Model")
]
# Generator Network Operations
self.gen_update_ops = tf.get_collection(
tf.GraphKeys.UPDATE_OPS, scope="FenceGAN/Generator_Model")
# Discriminator Network Operations
self.disc_update_ops = tf.get_collection(
tf.GraphKeys.UPDATE_OPS, scope="FenceGAN/Discriminator_Model")
with tf.control_dependencies(self.gen_update_ops):
self.gen_op = self.generator_optimizer.minimize(
self.loss_generator,
var_list=self.generator_vars,
global_step=self.global_step_tensor,
)
with tf.control_dependencies(self.disc_update_ops):
self.disc_op = self.discriminator_optimizer.minimize(
self.loss_discriminator, var_list=self.discriminator_vars)
# Exponential Moving Average for Estimation
self.dis_ema = tf.train.ExponentialMovingAverage(
decay=self.config.trainer.ema_decay)
maintain_averages_op_dis = self.dis_ema.apply(
self.discriminator_vars)
self.gen_ema = tf.train.ExponentialMovingAverage(
decay=self.config.trainer.ema_decay)
maintain_averages_op_gen = self.gen_ema.apply(self.generator_vars)
with tf.control_dependencies([self.disc_op]):
self.train_dis_op = tf.group(maintain_averages_op_dis)
with tf.control_dependencies([self.gen_op]):
self.train_gen_op = tf.group(maintain_averages_op_gen)
# Build Test Graph
self.logger.info("Building Testing Graph...")
with tf.variable_scope("FenceGAN"):
with tf.variable_scope("Generator_Model"):
self.image_gen_ema = self.generator(self.noise_tensor,
getter=get_getter(
self.gen_ema))
with tf.variable_scope("Discriminator_Model"):
self.disc_real_ema, self.f_layer_real_ema = self.discriminator(
self.image_input,
do_spectral_norm=self.config.trainer.do_spectral_norm,
getter=get_getter(self.dis_ema),
)
self.disc_fake_ema, self.f_layer_fake_ema = self.discriminator(
self.image_gen,
do_spectral_norm=self.config.trainer.do_spectral_norm,
getter=get_getter(self.dis_ema),
)
with tf.name_scope("Testing"):
with tf.name_scope("Image_Based"):
delta = self.image_input - self.image_gen_ema
delta_flat = tf.layers.Flatten()(delta)
img_score_l1 = tf.norm(delta_flat,
ord=1,
axis=1,
keepdims=False,
name="img_loss__1")
self.img_score_l1 = tf.squeeze(img_score_l1)
delta = self.image_input - self.image_gen_ema
delta_flat = tf.layers.Flatten()(delta)
img_score_l2 = tf.norm(delta_flat,
ord=2,
axis=1,
keepdims=False,
name="img_loss__2")
self.img_score_l2 = tf.squeeze(img_score_l2)
with tf.name_scope("Noise_Based"):
delta = self.embedding_rec_ema - self.embedding_q_ema
delta_flat = tf.layers.Flatten()(delta)
z_score_l1 = tf.norm(delta_flat,
ord=1,
axis=1,
keepdims=False,
name="z_loss_1")
self.z_score_l1 = tf.squeeze(z_score_l1)
delta = self.embedding_rec_ema - self.embedding_q_ema
delta_flat = tf.layers.Flatten()(delta)
z_score_l2 = tf.norm(delta_flat,
ord=2,
axis=1,
keepdims=False,
name="z_loss_2")
self.z_score_l2 = tf.squeeze(z_score_l2)
# Tensorboard
if self.config.log.enable_summary:
with tf.name_scope("train_summary"):
with tf.name_scope("dis_summary"):
tf.summary.scalar("loss_disc", self.loss_discriminator,
["dis"])
tf.summary.scalar("loss_disc_real", self.disc_loss_real,
["dis"])
tf.summary.scalar("loss_disc_fake", self.disc_loss_fake,
["dis"])
with tf.name_scope("gen_summary"):
tf.summary.scalar("loss_generator", self.loss_generator,
["gen"])
with tf.name_scope("img_summary"):
tf.summary.image("input_image", self.image_input, 3,
["img"])
tf.summary.image("reconstructed", self.image_gen, 3,
["img"])
self.sum_op_dis = tf.summary.merge_all("dis")
self.sum_op_gen = tf.summary.merge_all("gen")
self.sum_op = tf.summary.merge([self.sum_op_dis, self.sum_op_gen])
self.sum_op_im = tf.summary.merge_all("img")
def build_model(self):
# Kernel initialization for the convolutions
self.init_kernel = tf.random_normal_initializer(mean=0.0, stddev=0.02)
# Placeholders
self.is_training = tf.placeholder(tf.bool)
self.image_input = tf.placeholder(tf.float32,
shape=[None] +
self.config.trainer.image_dims,
name="x")
self.noise_tensor = tf.placeholder(
tf.float32,
shape=[None, self.config.trainer.noise_dim],
name="noise")
self.true_labels = tf.placeholder(dtype=tf.float32,
shape=[None, 1],
name="true_labels")
self.generated_labels = tf.placeholder(dtype=tf.float32,
shape=[None, 1],
name="gen_labels")
self.real_noise = tf.placeholder(dtype=tf.float32,
shape=[None] +
self.config.trainer.image_dims,
name="real_noise")
self.fake_noise = tf.placeholder(dtype=tf.float32,
shape=[None] +
self.config.trainer.image_dims,
name="fake_noise")
self.logger.info("Building training graph...")
with tf.variable_scope("BIGAN"):
with tf.variable_scope("Encoder_Model"):
self.noise_gen = self.encoder(self.image_input)
with tf.variable_scope("Generator_Model"):
self.image_gen = self.generator(
self.noise_tensor) + self.fake_noise
self.reconstructed = self.generator(self.noise_gen)
with tf.variable_scope("Discriminator_Model"):
# E(x) and x --> This being real is the output of discriminator
l_encoder, inter_layer_inp = self.discriminator(
self.noise_gen, self.image_input + self.real_noise)
# z and G(z)
l_generator, inter_layer_rct = self.discriminator(
self.noise_tensor, self.image_gen)
# Loss Function Implementations
with tf.name_scope("Loss_Functions"):
# Discriminator
# Discriminator sees the encoder result as true because it discriminates E(x), x as the real pair
self.loss_dis_enc = tf.reduce_mean(
tf.nn.sigmoid_cross_entropy_with_logits(
labels=self.true_labels, logits=l_encoder))
self.loss_dis_gen = tf.reduce_mean(
tf.nn.sigmoid_cross_entropy_with_logits(
labels=self.generated_labels, logits=l_generator))
self.loss_discriminator = self.loss_dis_enc + self.loss_dis_gen
if self.config.trainer.flip_labels:
labels_gen = tf.zeros_like(l_generator)
labels_enc = tf.ones_like(l_encoder)
else:
labels_gen = tf.ones_like(l_generator)
labels_enc = tf.zeros_like(l_encoder)
# Generator
# Generator is considered as the true ones here because it tries to fool discriminator
self.loss_generator = tf.reduce_mean(
tf.nn.sigmoid_cross_entropy_with_logits(labels=labels_gen,
logits=l_generator))
# Encoder
# Encoder is considered as the fake one because it tries to fool the discriminator also
self.loss_encoder = tf.reduce_mean(
tf.nn.sigmoid_cross_entropy_with_logits(labels=labels_enc,
logits=l_encoder))
# Optimizer Implementations
with tf.name_scope("Optimizers"):
# Build the optimizers
self.generator_optimizer = tf.train.AdamOptimizer(
self.config.trainer.generator_l_rate,
beta1=self.config.trainer.optimizer_adam_beta1,
beta2=self.config.trainer.optimizer_adam_beta2,
)
self.discriminator_optimizer = tf.train.AdamOptimizer(
self.config.trainer.discriminator_l_rate,
beta1=self.config.trainer.optimizer_adam_beta1,
beta2=self.config.trainer.optimizer_adam_beta2,
)
self.encoder_optimizer = tf.train.AdamOptimizer(
self.config.trainer.generator_l_rate,
beta1=self.config.trainer.optimizer_adam_beta1,
beta2=self.config.trainer.optimizer_adam_beta2,
)
# Collect all the variables
all_variables = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES)
# Generator Network Variables
self.generator_vars = [
v for v in all_variables
if v.name.startswith("BIGAN/Generator_Model")
]
# Discriminator Network Variables
self.discriminator_vars = [
v for v in all_variables
if v.name.startswith("BIGAN/Discriminator_Model")
]
# Encoder Network Variables
self.encoder_vars = [
v for v in all_variables
if v.name.startswith("BIGAN/Encoder_Model")
]
# Create Training Operations
# Generator Network Operations
self.gen_update_ops = tf.get_collection(
tf.GraphKeys.UPDATE_OPS, scope="BIGAN/Generator_Model")
# Discriminator Network Operations
self.disc_update_ops = tf.get_collection(
tf.GraphKeys.UPDATE_OPS, scope="BIGAN/Discriminator_Model")
# Encoder Network Operations
self.enc_update_ops = tf.get_collection(
tf.GraphKeys.UPDATE_OPS, scope="BIGAN/Encoder_Model")
# Initialization of Optimizers
with tf.control_dependencies(self.gen_update_ops):
self.gen_op = self.generator_optimizer.minimize(
self.loss_generator,
var_list=self.generator_vars,
global_step=self.global_step_tensor,
)
with tf.control_dependencies(self.disc_update_ops):
self.disc_op = self.discriminator_optimizer.minimize(
self.loss_discriminator, var_list=self.discriminator_vars)
with tf.control_dependencies(self.enc_update_ops):
self.enc_op = self.encoder_optimizer.minimize(
self.loss_encoder, var_list=self.encoder_vars)
# Exponential Moving Average for Estimation
self.dis_ema = tf.train.ExponentialMovingAverage(
decay=self.config.trainer.ema_decay)
maintain_averages_op_dis = self.dis_ema.apply(
self.discriminator_vars)
self.gen_ema = tf.train.ExponentialMovingAverage(
decay=self.config.trainer.ema_decay)
maintain_averages_op_gen = self.gen_ema.apply(self.generator_vars)
self.enc_ema = tf.train.ExponentialMovingAverage(
decay=self.config.trainer.ema_decay)
maintain_averages_op_enc = self.enc_ema.apply(self.encoder_vars)
with tf.control_dependencies([self.disc_op]):
self.train_dis_op = tf.group(maintain_averages_op_dis)
with tf.control_dependencies([self.gen_op]):
self.train_gen_op = tf.group(maintain_averages_op_gen)
with tf.control_dependencies([self.enc_op]):
self.train_enc_op = tf.group(maintain_averages_op_enc)
self.logger.info("Building Testing Graph...")
with tf.variable_scope("BIGAN"):
with tf.variable_scope("Encoder_Model"):
self.noise_gen_ema = self.encoder(self.image_input,
getter=get_getter(
self.enc_ema))
with tf.variable_scope("Generator_Model"):
self.reconstruct_ema = self.generator(self.noise_gen_ema,
getter=get_getter(
self.gen_ema))
with tf.variable_scope("Discriminator_Model"):
self.l_encoder_ema, self.inter_layer_inp_ema = self.discriminator(
self.noise_gen_ema, # E(x)
self.image_input, # x
getter=get_getter(self.dis_ema),
)
self.l_generator_ema, self.inter_layer_rct_ema = self.discriminator(
self.noise_gen_ema, # E(x)
self.reconstruct_ema, # G(E(x))
getter=get_getter(self.dis_ema),
)
with tf.name_scope("Testing"):
with tf.variable_scope("Reconstruction_Loss"):
# LG(x) = ||x - G(E(x))||_1
delta = self.image_input - self.reconstruct_ema
delta_flat = tf.layers.Flatten()(delta)
self.gen_score = tf.norm(
delta_flat,
ord=self.config.trainer.degree,
axis=1,
keepdims=False,
name="epsilon",
)
with tf.variable_scope("Discriminator_Loss"):
if self.config.trainer.loss_method == "cross_e":
self.dis_score = tf.nn.sigmoid_cross_entropy_with_logits(
labels=tf.ones_like(self.l_encoder_ema),
logits=self.l_encoder_ema)
elif self.config.trainer.loss_method == "fm":
fm = self.inter_layer_inp_ema - self.inter_layer_rct_ema
fm = tf.layers.Flatten()(fm)
self.dis_score = tf.norm(fm,
ord=self.config.trainer.degree,
axis=1,
keepdims=False,
name="d_loss")
self.dis_score = tf.squeeze(self.dis_score)
with tf.variable_scope("Score"):
self.list_scores = (
1 - self.config.trainer.weight
) * self.dis_score + self.config.trainer.weight * self.gen_score
if self.config.trainer.enable_early_stop:
self.rec_error_valid = tf.reduce_mean(self.list_scores)
if self.config.log.enable_summary:
with tf.name_scope("Summary"):
with tf.name_scope("Disc_Summary"):
tf.summary.scalar("loss_discriminator",
self.loss_discriminator, ["dis"])
tf.summary.scalar("loss_dis_encoder", self.loss_dis_enc,
["dis"])
tf.summary.scalar("loss_dis_gen", self.loss_dis_gen,
["dis"])
with tf.name_scope("Gen_Summary"):
tf.summary.scalar("loss_generator", self.loss_generator,
["gen"])
tf.summary.scalar("loss_encoder", self.loss_encoder,
["gen"])
with tf.name_scope("Image_Summary"):
tf.summary.image("reconstruct", self.reconstructed, 3,
["image"])
tf.summary.image("input_images", self.image_input, 3,
["image"])
if self.config.trainer.enable_early_stop:
with tf.name_scope("validation_summary"):
tf.summary.scalar("valid", self.rec_error_valid, ["v"])
self.sum_op_dis = tf.summary.merge_all("dis")
self.sum_op_gen = tf.summary.merge_all("gen")
self.sum_op_im = tf.summary.merge_all("image")
self.sum_op_valid = tf.summary.merge_all("v")
def build_model(self):
# Kernel initialization for the convolutions
self.init_kernel = tf.random_normal_initializer(mean=0.0, stddev=0.02)
# Placeholders
self.is_training = tf.placeholder(tf.bool)
self.image_input = tf.placeholder(tf.float32,
shape=[None] +
self.config.trainer.image_dims,
name="x")
self.true_labels = tf.placeholder(dtype=tf.float32,
shape=[None, 1],
name="true_labels")
self.generated_labels = tf.placeholder(dtype=tf.float32,
shape=[None, 1],
name="gen_labels")
self.logger.info("Building training graph...")
with tf.variable_scope("Mark1"):
with tf.variable_scope("Generator_Model"):
self.noise_gen, self.img_rec, self.noise_rec = self.generator(
self.image_input)
# Discriminator results of (G(z),z) and (x, E(x))
with tf.variable_scope("Discriminator_Model_XZ"):
l_generator, inter_layer_rct_xz = self.discriminator_xz(
self.img_rec,
self.noise_gen,
do_spectral_norm=self.config.spectral_norm)
l_encoder, inter_layer_inp_xz = self.discriminator_xz(
self.image_input,
self.noise_gen,
do_spectral_norm=self.config.do_spectral_norm)
# Discrimeinator results of (x, x) and (x, G(E(x))
with tf.variable_scope("Discriminator_Model_XX"):
x_logit_real, inter_layer_inp_xx = self.discriminator_xx(
self.image_input,
self.image_input,
do_spectral_norm=self.config.spectral_norm)
x_logit_fake, inter_layer_rct_xx = self.discriminator_xx(
self.image_input,
self.img_rec,
do_spectral_norm=self.config.spectral_norm)
# Discriminator results of (z, z) and (z, E(G(z))
with tf.variable_scope("Discriminator_Model_ZZ"):
z_logit_real, _ = self.discriminator_zz(
self.noise_gen,
self.noise_gen,
do_spectral_norm=self.config.spectral_norm)
z_logit_fake, _ = self.discriminator_zz(
self.noise_gen,
self.noise_rec,
do_spectral_norm=self.config.spectral_norm)
with tf.name_scope("Loss_Functions"):
# discriminator xz
# Discriminator should classify encoder pair as real
loss_dis_enc = tf.reduce_mean(
tf.nn.sigmoid_cross_entropy_with_logits(
labels=self.true_labels, logits=l_encoder))
# Discriminator should classify generator pair as fake
loss_dis_gen = tf.reduce_mean(
tf.nn.sigmoid_cross_entropy_with_logits(
labels=self.generated_labels, logits=l_generator))
self.dis_loss_xz = loss_dis_gen + loss_dis_enc
# discriminator xx
x_real_dis = tf.nn.sigmoid_cross_entropy_with_logits(
logits=x_logit_real, labels=self.true_labels)
x_fake_dis = tf.nn.sigmoid_cross_entropy_with_logits(
logits=x_logit_fake, labels=self.generated_labels)
self.dis_loss_xx = tf.reduce_mean(x_real_dis + x_fake_dis)
# discriminator zz
z_real_dis = tf.nn.sigmoid_cross_entropy_with_logits(
logits=z_logit_real, labels=self.true_labels)
z_fake_dis = tf.nn.sigmoid_cross_entropy_with_logits(
logits=z_logit_fake, labels=self.generated_labels)
self.dis_loss_zz = tf.reduce_mean(z_real_dis + z_fake_dis)
# Feature matching part
fm = inter_layer_inp_xx - inter_layer_rct_xx
fm = tf.layers.Flatten()(fm)
self.feature_match = tf.reduce_mean(
tf.norm(fm, ord=2, axis=1, keepdims=False))
# Compute the whole discriminator loss
self.loss_discriminator = (
self.dis_loss_xz + self.dis_loss_xx + self.dis_loss_zz
if self.config.trainer.allow_zz else self.dis_loss_xz +
self.dis_loss_xx + self.feature_match if
self.config.trainer.loss_method == "fm" else self.dis_loss_xz +
self.dis_loss_xx)
# Generator
# Adversarial Loss
if self.config.trainer.flip_labels:
labels_gen = tf.zeros_like(l_generator)
labels_enc = tf.ones_like(l_encoder)
else:
labels_gen = tf.ones_like(l_generator)
labels_enc = tf.zeros_like(l_encoder)
self.gen_loss_enc = tf.reduce_mean(
tf.nn.sigmoid_cross_entropy_with_logits(labels=labels_enc,
logits=l_encoder))
self.gen_loss_gen = tf.reduce_mean(
tf.nn.sigmoid_cross_entropy_with_logits(labels=labels_gen,
logits=l_generator))
# Contextual Loss
l1_norm = self.image_input - self.img_rec
l1_norm = tf.layers.Flatten()(l1_norm)
self.gen_loss_con = tf.reduce_mean(
tf.norm(l1_norm, ord=1, axis=1, keepdims=False))
# Encoder Loss
l2_norm = self.noise_gen - self.noise_rec
l2_norm = tf.layers.Flatten()(l2_norm)
self.gen_loss_enc = tf.reduce_mean(
tf.norm(l2_norm, ord=2, axis=1, keepdims=False))
self.gen_loss_ce = self.gen_loss_enc + self.gen_loss_gen
self.gen_loss_total = (
self.config.trainer.weight_adv * self.gen_loss_ce +
self.config.trainer.weight_cont * self.gen_loss_con +
self.config.trainer.weight_enc * self.gen_loss_enc)
with tf.name_scope("Optimizers"):
# Collect all the variables
all_variables = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES)
# Generator Network Variables
self.gvars = [
v for v in all_variables
if v.name.startswith("Mark1/Generator_Model")
]
self.dxzvars = [
v for v in all_variables
if v.name.startswith("Mark1/Discriminator_Model_XZ")
]
self.dxxvars = [
v for v in all_variables
if v.name.startswith("Mark1/Discriminator_Model_XX")
]
self.dzzvars = [
v for v in all_variables
if v.name.startswith("Mark1/Discriminator_Model_ZZ")
]
# Create Training Operations
# Generator Network Operations
self.update_ops_gen = tf.get_collection(
tf.GraphKeys.UPDATE_OPS, scope="Mark1/Generator_Model")
self.update_ops_enc = tf.get_collection(
tf.GraphKeys.UPDATE_OPS, scope="Mark1/Encoder_Model")
self.update_ops_dis_xz = tf.get_collection(
tf.GraphKeys.UPDATE_OPS, scope="Mark1/Discriminator_Model_XZ")
self.update_ops_dis_xx = tf.get_collection(
tf.GraphKeys.UPDATE_OPS, scope="Mark1/Discriminator_Model_XX")
self.update_ops_dis_zz = tf.get_collection(
tf.GraphKeys.UPDATE_OPS, scope="Mark1/Discriminator_Model_ZZ")
self.disc_optimizer = tf.train.AdamOptimizer(
learning_rate=self.config.trainer.discriminator_l_rate,
beta1=self.config.trainer.optimizer_adam_beta1,
beta2=self.config.trainer.optimizer_adam_beta2,
)
self.gen_optimizer = tf.train.AdamOptimizer(
learning_rate=self.config.trainer.generator_l_rate,
beta1=self.config.trainer.optimizer_adam_beta1,
beta2=self.config.trainer.optimizer_adam_beta2,
)
# Initialization of Optimizers
with tf.control_dependencies(self.update_ops_gen):
self.gen_op = self.gen_optimizer.minimize(
self.gen_loss_total,
global_step=self.global_step_tensor,
var_list=self.gvars)
with tf.control_dependencies(self.update_ops_dis_xz):
self.dis_op_xz = self.disc_optimizer.minimize(
self.dis_loss_xz, var_list=self.dxzvars)
with tf.control_dependencies(self.update_ops_dis_xx):
self.dis_op_xx = self.disc_optimizer.minimize(
self.dis_loss_xx, var_list=self.dxxvars)
with tf.control_dependencies(self.update_ops_dis_zz):
self.dis_op_zz = self.disc_optimizer.minimize(
self.dis_loss_zz, var_list=self.dzzvars)
# Exponential Moving Average for inference
def train_op_with_ema_dependency(vars, op):
ema = tf.train.ExponentialMovingAverage(
decay=self.config.trainer.ema_decay)
maintain_averages_op = ema.apply(vars)
with tf.control_dependencies([op]):
train_op = tf.group(maintain_averages_op)
return train_op, ema
self.train_gen_op, self.gen_ema = train_op_with_ema_dependency(
self.gvars, self.gen_op)
self.train_dis_op_xz, self.xz_ema = train_op_with_ema_dependency(
self.dxzvars, self.dis_op_xz)
self.train_dis_op_xx, self.xx_ema = train_op_with_ema_dependency(
self.dxxvars, self.dis_op_xx)
self.train_dis_op_zz, self.zz_ema = train_op_with_ema_dependency(
self.dzzvars, self.dis_op_zz)
self.logger.info("Building Testing Graph...")
with tf.variable_scope("Mark1"):
with tf.variable_scope("Generator_Model"):
self.noise_gen_ema, self.img_rec_ema, self.noise_rec_ema = self.generator(
self.image_input, getter=get_getter(self.gen_ema))
with tf.name_scope("Testing"):
with tf.variable_scope("Reconstruction_Loss"):
# | G_E(x) - E(G(x))|1
# Difference between the noise generated from the input image and reconstructed noise
delta = self.noise_gen_ema - self.noise_rec_ema
delta = tf.layers.Flatten()(delta)
self.score = tf.norm(delta, ord=1, axis=1, keepdims=False)
if self.config.trainer.enable_early_stop:
self.rec_error_valid = tf.reduce_mean(self.score)
if self.config.log.enable_summary:
with tf.name_scope("summary"):
with tf.name_scope("disc_summary"):
tf.summary.scalar("loss_discriminator",
self.loss_discriminator, ["dis"])
tf.summary.scalar("loss_dis_encoder", loss_dis_enc,
["dis"])
tf.summary.scalar("loss_dis_gen", loss_dis_gen, ["dis"])
tf.summary.scalar("loss_dis_xz", self.dis_loss_xz, ["dis"])
tf.summary.scalar("loss_dis_xx", self.dis_loss_xx, ["dis"])
if self.config.trainer.allow_zz:
tf.summary.scalar("loss_dis_zz", self.dis_loss_zz,
["dis"])
if self.config.trainer.loss_method:
tf.summary.scalar("loss_dis_fm", self.feature_match,
["dis"])
with tf.name_scope("gen_summary"):
tf.summary.scalar("loss_generator_total",
self.gen_loss_total, ["gen"])
tf.summary.scalar("loss_gen_adv", self.gen_loss_ce,
["gen"])
tf.summary.scalar("loss_gen_con", self.gen_loss_con,
["gen"])
tf.summary.scalar("loss_gen_enc", self.gen_loss_enc,
["gen"])
with tf.name_scope("image_summary"):
tf.summary.image("reconstruct", self.img_rec, 3, ["image"])
tf.summary.image("input_images", self.image_input, 3,
["image"])
if self.config.trainer.enable_early_stop:
with tf.name_scope("validation_summary"):
tf.summary.scalar("valid", self.rec_error_valid, ["v"])
self.sum_op_dis = tf.summary.merge_all("dis")
self.sum_op_gen = tf.summary.merge_all("gen")
self.sum_op_im = tf.summary.merge_all("image")
self.sum_op_valid = tf.summary.merge_all("v")
