def bulid_mrgan(self): # Generator self.g = self.generator(self.z) self.g_reg = self.generator(self.encoder(self.x), reuse=True) # Discriminator d_real = self.discriminator(self.x) d_real_reg = self.discriminator(self.g_reg, reuse=True) d_fake = self.discriminator(self.g, reuse=True) # Losses # Manifold Step # d_loss_1 = tf.reduce_mean(t.safe_log(d_real) + t.safe_log(1. - d_real_reg)) # g_loss_1 = tf.reduce_mean(self.lambda_1 * t.safe_log(d_real_reg)) - \ # t.mse_loss(self.x, self.g_reg, self.batch_size) # Diffusion Step # d_loss_2 = tf.reduce_mean(t.safe_log(d_real_reg) + t.safe_log(1. - d_fake)) # g_loss_2 = tf.reduce_mean(t.safe_log(d_fake)) d_real_loss = -tf.reduce_mean(t.safe_log(d_real)) d_fake_loss = -tf.reduce_mean(t.safe_log(1. - d_fake)) self.d_loss = d_real_loss + d_fake_loss e_mse_loss = self.lambda_1 * t.mse_loss( self.x, self.g_reg, self.batch_size, is_mean=True) e_adv_loss = self.lambda_2 * tf.reduce_mean(t.safe_log(d_real_reg)) self.e_loss = e_adv_loss + e_mse_loss self.g_loss = -tf.reduce_mean(t.safe_log(d_fake)) + self.e_loss # Summary tf.summary.scalar("loss/d_real_loss", d_real_loss) tf.summary.scalar("loss/d_fake_loss", d_fake_loss) tf.summary.scalar("loss/d_loss", self.d_loss) tf.summary.scalar("loss/e_adv_loss", e_adv_loss) tf.summary.scalar("loss/e_mse_loss", e_mse_loss) tf.summary.scalar("loss/e_loss", self.e_loss) tf.summary.scalar("loss/g_loss", self.g_loss) # Collect trainer values t_vars = tf.trainable_variables() d_params = [v for v in t_vars if v.name.startswith('d')] g_params = [v for v in t_vars if v.name.startswith('g')] e_params = [v for v in t_vars if v.name.startswith('e')] # Optimizer self.d_op = tf.train.AdamOptimizer(learning_rate=self.lr, beta1=self.beta1).minimize( self.d_loss, var_list=d_params) self.g_op = tf.train.AdamOptimizer(learning_rate=self.lr, beta1=self.beta1).minimize( self.g_loss, var_list=g_params) self.e_op = tf.train.AdamOptimizer(learning_rate=self.lr, beta1=self.beta1).minimize( self.e_loss, var_list=e_params) # Merge summary self.merged = tf.summary.merge_all() # Model Saver self.saver = tf.train.Saver(max_to_keep=1) self.writer = tf.summary.FileWriter('./model/', self.s.graph)
def build_ebgan(self): # Generator self.g = self.generator(self.z) self.g_test = self.generator(self.z, reuse=True, is_train=False) # Discriminator d_embed_real, d_decode_real = self.discriminator(self.x) d_embed_fake, d_decode_fake = self.discriminator(self.g, reuse=True) d_real_loss = t.mse_loss(d_decode_real, self.x, self.batch_size) d_fake_loss = t.mse_loss(d_decode_fake, self.g, self.batch_size) self.d_loss = d_real_loss + tf.maximum(0., self.margin - d_fake_loss) if self.EnablePullAway: self.pt_loss = self.pullaway_loss(d_embed_fake, self.batch_size) self.g_loss = d_fake_loss + self.pt_lambda * self.pt_loss # Summary tf.summary.scalar("loss/d_real_loss", d_real_loss) tf.summary.scalar("loss/d_fake_loss", d_fake_loss) tf.summary.scalar("loss/d_loss", self.d_loss) tf.summary.scalar("loss/g_loss", self.g_loss) tf.summary.scalar("loss/pt_loss", self.pt_loss) # Optimizer t_vars = tf.trainable_variables() d_params = [v for v in t_vars if v.name.startswith('d')] g_params = [v for v in t_vars if v.name.startswith('g')] self.d_op = tf.train.AdamOptimizer(learning_rate=self.d_lr, beta1=self.beta1, beta2=self.beta2).minimize( self.d_loss, var_list=d_params) self.g_op = tf.train.AdamOptimizer(learning_rate=self.g_lr, beta1=self.beta1, beta2=self.beta2).minimize( self.g_loss, var_list=g_params) # Merge summary self.merged = tf.summary.merge_all() # Model saver self.saver = tf.train.Saver(max_to_keep=1) self.writer = tf.summary.FileWriter('./model/', self.s.graph)
def build_lsgan(self): # Generator self.g = self.generator(self.z) # Discriminator d_real = self.discriminator(self.x) d_fake = self.discriminator(self.g, reuse=True) # LSGAN Loss d_real_loss = t.mse_loss(d_real, tf.ones_like(d_real), self.batch_size) d_fake_loss = t.mse_loss(d_fake, tf.zeros_like(d_fake), self.batch_size) self.d_loss = (d_real_loss + d_fake_loss) / 2. self.g_loss = t.mse_loss(d_fake, tf.ones_like(d_fake), self.batch_size) # Summary tf.summary.scalar("loss/d_real_loss", d_real_loss) tf.summary.scalar("loss/d_fake_loss", d_fake_loss) tf.summary.scalar("loss/d_loss", self.d_loss) tf.summary.scalar("loss/g_loss", self.g_loss) # optimizer t_vars = tf.trainable_variables() d_params = [v for v in t_vars if v.name.startswith('d')] g_params = [v for v in t_vars if v.name.startswith('g')] self.d_op = tf.train.AdamOptimizer(learning_rate=self.lr, beta1=self.beta1).minimize( self.d_loss, var_list=d_params) self.g_op = tf.train.AdamOptimizer(learning_rate=self.lr, beta1=self.beta1).minimize( self.g_loss, var_list=g_params) # Merge summary self.merged = tf.summary.merge_all() # Model saver self.saver = tf.train.Saver(max_to_keep=1) self.writer = tf.summary.FileWriter('./model/', self.s.graph)
def build_magan(self): # Generator self.g = self.generator(self.z) # Discriminator _, d_real = self.discriminator(self.x) _, d_fake = self.discriminator(self.g, reuse=True) self.d_real_loss = t.mse_loss(self.x, d_real, self.batch_size) self.d_fake_loss = t.mse_loss(self.g, d_fake, self.batch_size) self.d_loss = self.d_real_loss + tf.maximum(0., self.m - self.d_fake_loss) self.g_loss = self.d_fake_loss # Summary tf.summary.scalar("loss/d_loss", self.d_loss) tf.summary.scalar("loss/d_real_loss", self.d_real_loss) tf.summary.scalar("loss/d_fake_loss", self.d_fake_loss) tf.summary.scalar("loss/g_loss", self.g_loss) # Optimizer t_vars = tf.trainable_variables() d_params = [v for v in t_vars if v.name.startswith('d')] g_params = [v for v in t_vars if v.name.startswith('g')] self.d_op = AdamaxOptimizer(learning_rate=self.lr, beta1=self.beta1).minimize( self.d_loss, var_list=d_params) self.g_op = AdamaxOptimizer(learning_rate=self.lr, beta1=self.beta1).minimize( self.g_loss, var_list=g_params) # Merge summary self.merged = tf.summary.merge_all() # Model saver self.saver = tf.train.Saver(max_to_keep=1) self.writer = tf.summary.FileWriter('./model/', self.s.graph)
def build_srgan(self): # Generator self.g = self.generator(self.x_lr) # Discriminator d_real = self.discriminator(self.x_hr) d_fake = self.discriminator(self.g, reuse=True) # Losses # d_real_loss = -tf.reduce_mean(t.safe_log(d_real)) # d_fake_loss = -tf.reduce_mean(t.safe_log(1. - d_fake)) d_real_loss = t.sce_loss(d_real, tf.ones_like(d_real)) d_fake_loss = t.sce_loss(d_fake, tf.zeros_like(d_fake)) self.d_loss = d_real_loss + d_fake_loss if self.use_vgg19: x_vgg_real = tf.image.resize_images(self.x_hr, size=self.vgg_image_shape[:2], align_corners=False) x_vgg_fake = tf.image.resize_images(self.g, size=self.vgg_image_shape[:2], align_corners=False) vgg_bottle_real = self.build_vgg19(x_vgg_real) vgg_bottle_fake = self.build_vgg19(x_vgg_fake, reuse=True) self.g_cnt_loss = self.cnt_scaling * t.mse_loss(vgg_bottle_fake, vgg_bottle_real, self.batch_size, is_mean=True) else: self.g_cnt_loss = t.mse_loss(self.g, self.x_hr, self.batch_size, is_mean=True) # self.g_adv_loss = self.adv_scaling * tf.reduce_mean(-1. * t.safe_log(d_fake)) self.g_adv_loss = self.adv_scaling * t.sce_loss( d_fake, tf.ones_like(d_fake)) self.g_loss = self.g_adv_loss + self.g_cnt_loss def inverse_transform(img): return (img + 1.) * 127.5 # calculate PSNR g, x_hr = inverse_transform(self.g), inverse_transform(self.x_hr) self.psnr = t.psnr_loss(g, x_hr, self.batch_size) # Summary tf.summary.scalar("loss/d_real_loss", d_real_loss) tf.summary.scalar("loss/d_fake_loss", d_fake_loss) tf.summary.scalar("loss/d_loss", self.d_loss) tf.summary.scalar("loss/g_cnt_loss", self.g_cnt_loss) tf.summary.scalar("loss/g_adv_loss", self.g_adv_loss) tf.summary.scalar("loss/g_loss", self.g_loss) tf.summary.scalar("misc/psnr", self.psnr) tf.summary.scalar("misc/lr", self.lr) # Optimizer t_vars = tf.trainable_variables() d_params = [v for v in t_vars if v.name.startswith('d')] g_params = [v for v in t_vars if v.name.startswith('g')] self.d_op = tf.train.AdamOptimizer(learning_rate=self.lr, beta1=self.beta1, beta2=self.beta2).minimize( loss=self.d_loss, var_list=d_params) self.g_op = tf.train.AdamOptimizer(learning_rate=self.lr, beta1=self.beta1, beta2=self.beta2).minimize( loss=self.g_loss, var_list=g_params) # pre-train self.g_init_op = tf.train.AdamOptimizer(learning_rate=self.lr, beta1=self.beta1, beta2=self.beta2).minimize( loss=self.g_cnt_loss, var_list=g_params) # Merge summary self.merged = tf.summary.merge_all() # Model saver self.saver = tf.train.Saver(max_to_keep=2) self.writer = tf.summary.FileWriter('./model/', self.s.graph)