From 93a930fd52e1520eaa20b0c76d8d8b3b2cab951d Mon Sep 17 00:00:00 2001 From: TalosOfCrete Date: Mon, 8 Jun 2020 17:58:16 -0500 Subject: [PATCH] upd --- core/leras/optimizers/RMSprop.py | 9 ++++++--- models/Model_SAEHD/Model.py | 27 +++++++++++++++++++-------- 2 files changed, 25 insertions(+), 11 deletions(-) diff --git a/core/leras/optimizers/RMSprop.py b/core/leras/optimizers/RMSprop.py index edd4c38..f7b948e 100644 --- a/core/leras/optimizers/RMSprop.py +++ b/core/leras/optimizers/RMSprop.py @@ -25,7 +25,7 @@ class RMSprop(nn.OptimizerBase): def get_weights(self): return [self.lr, self.rho, self.epsilon, self.iterations] + list(self.accumulators_dict.values()) - def initialize_variables(self, trainable_weights, vars_on_cpu=True): + def initialize_variables(self, trainable_weights, vars_on_cpu=True, lr_dropout_on_cpu=False): # Initialize here all trainable variables used in training e = tf.device('/CPU:0') if vars_on_cpu else None if e: e.__enter__() @@ -34,8 +34,11 @@ class RMSprop(nn.OptimizerBase): self.accumulators_dict.update ( accumulators) if self.lr_dropout != 1.0: + e = tf.device('/CPU:0') if lr_dropout_on_cpu else None + if e: e.__enter__() lr_rnds = [ nn.random_binomial( v.shape, p=self.lr_dropout, dtype=v.dtype) for v in trainable_weights ] - self.lr_rnds_dict.update ( { v.name : rnd for v,rnd in zip(trainable_weights,lr_rnds) } ) + if e: e.__exit__(None, None, None) + elf.lr_rnds_dict.update ( { v.name : rnd for v,rnd in zip(trainable_weights,lr_rnds) } ) if e: e.__exit__(None, None, None) def get_update_op(self, grads_vars): @@ -66,4 +69,4 @@ class RMSprop(nn.OptimizerBase): updates.append (state_ops.assign(v, new_v)) return control_flow_ops.group ( *updates, name=self.name+'_updates') -nn.RMSprop = RMSprop \ No newline at end of file +nn.RMSprop = RMSprop diff --git a/models/Model_SAEHD/Model.py b/models/Model_SAEHD/Model.py index 959453b..418b1de 100644 --- a/models/Model_SAEHD/Model.py +++ b/models/Model_SAEHD/Model.py @@ -42,7 +42,8 @@ class SAEHDModel(ModelBase): default_d_mask_dims = self.options['d_mask_dims'] = self.options.get('d_mask_dims', None) default_masked_training = self.options['masked_training'] = self.load_or_def_option('masked_training', True) default_eyes_prio = self.options['eyes_prio'] = self.load_or_def_option('eyes_prio', False) - default_lr_dropout = self.options['lr_dropout'] = self.load_or_def_option('lr_dropout', False) + default_lr_dropout = self.options['lr_dropout'] = self.load_or_def_option('lr_dropout', 'n') + default_lr_dropout = {True:'y', False:'n'}.get(default_lr_dropout, default_lr_dropout) #backward comp default_random_warp = self.options['random_warp'] = self.load_or_def_option('random_warp', True) default_gan_power = self.options['gan_power'] = self.load_or_def_option('gan_power', 0.0) default_true_face_power = self.options['true_face_power'] = self.load_or_def_option('true_face_power', 0.0) @@ -106,7 +107,8 @@ class SAEHDModel(ModelBase): if self.is_first_run() or ask_override: self.options['models_opt_on_gpu'] = io.input_bool ("Place models and optimizer on GPU", default_models_opt_on_gpu, help_message="When you train on one GPU, by default model and optimizer weights are placed on GPU to accelerate the process. You can place they on CPU to free up extra VRAM, thus set bigger dimensions.") - self.options['lr_dropout'] = io.input_bool ("Use learning rate dropout", default_lr_dropout, help_message="When the face is trained enough, you can enable this option to get extra sharpness and reduce subpixel shake for fewer iterations.") + self.options['lr_dropout'] = io.input_str (f"Use learning rate dropout", default_lr_dropout, ['n','y','cpu'], help_message="When the face is trained enough, you can enable this option to get extra sharpness and reduce subpixel shake for fewer iterations.\nn - disabled.\ny - enabled\ncpu - enabled on CPU. This allows you to save VRAM, sacrificing 20% iteration speed.") + self.options['random_warp'] = io.input_bool ("Enable random warp of samples", default_random_warp, help_message="Random warp is required to generalize facial expressions of both faces. When the face is trained enough, you can disable it to get extra sharpness and reduce subpixel shake for fewer iterations.") self.options['gan_power'] = np.clip ( io.input_number ("GAN power", default_gan_power, add_info="0.0 .. 10.0", help_message="Train the network in Generative Adversarial manner. Accelerates the speed of training. Forces the neural network to learn small details of the face. You can enable/disable this option at any time. Typical value is 1.0"), 0.0, 10.0 ) @@ -238,7 +240,7 @@ class SAEHDModel(ModelBase): # Initialize optimizers lr=5e-5 - lr_dropout = 0.3 if self.options['lr_dropout'] and not self.pretrain else 1.0 + lr_dropout = 0.3 if self.options['lr_dropout'] in ['y','cpu'] and not self.pretrain else 1.0 clipnorm = 1.0 if self.options['clipgrad'] else 0.0 self.src_dst_opt = nn.RMSprop(lr=lr, lr_dropout=lr_dropout, clipnorm=clipnorm, name='src_dst_opt') self.model_filename_list += [ (self.src_dst_opt, 'src_dst_opt.npy') ] @@ -247,16 +249,17 @@ class SAEHDModel(ModelBase): elif 'liae' in archi: self.src_dst_trainable_weights = self.encoder.get_weights() + self.inter_AB.get_weights() + self.inter_B.get_weights() + self.decoder.get_weights() - self.src_dst_opt.initialize_variables (self.src_dst_trainable_weights, vars_on_cpu=optimizer_vars_on_cpu) + self.src_dst_opt.initialize_variables (self.src_dst_trainable_weights, vars_on_cpu=optimizer_vars_on_cpu, lr_dropout_on_cpu=self.options['lr_dropout']=='cpu') + if self.options['true_face_power'] != 0: self.D_code_opt = nn.RMSprop(lr=lr, lr_dropout=lr_dropout, clipnorm=clipnorm, name='D_code_opt') - self.D_code_opt.initialize_variables ( self.code_discriminator.get_weights(), vars_on_cpu=optimizer_vars_on_cpu) + self.D_code_opt.initialize_variables ( self.code_discriminator.get_weights(), vars_on_cpu=optimizer_vars_on_cpu, lr_dropout_on_cpu=self.options['lr_dropout']=='cpu') self.model_filename_list += [ (self.D_code_opt, 'D_code_opt.npy') ] if gan_power != 0: self.D_src_dst_opt = nn.RMSprop(lr=lr, lr_dropout=lr_dropout, clipnorm=clipnorm, name='D_src_dst_opt') - self.D_src_dst_opt.initialize_variables ( self.D_src.get_weights()+self.D_dst.get_weights(), vars_on_cpu=optimizer_vars_on_cpu) + self.D_src_dst_opt.initialize_variables ( self.D_src.get_weights()+self.D_dst.get_weights(), vars_on_cpu=optimizer_vars_on_cpu, lr_dropout_on_cpu=self.options['lr_dropout']=='cpu') self.model_filename_list += [ (self.D_src_dst_opt, 'D_src_dst_opt.npy') ] if self.is_training: @@ -343,7 +346,11 @@ class SAEHDModel(ModelBase): gpu_psd_target_dst_masked = gpu_pred_src_dst*gpu_target_dstm_blur gpu_psd_target_dst_anti_masked = gpu_pred_src_dst*(1.0 - gpu_target_dstm_blur) - gpu_src_loss = tf.reduce_mean ( 10*nn.dssim(gpu_target_src_masked_opt, gpu_pred_src_src_masked_opt, max_val=1.0, filter_size=int(resolution/11.6)), axis=[1]) + if resolution < 256: + gpu_src_loss = tf.reduce_mean ( 10*nn.dssim(gpu_target_src_masked_opt, gpu_pred_src_src_masked_opt, max_val=1.0, filter_size=int(resolution/11.6)), axis=[1]) + else: + gpu_src_loss = tf.reduce_mean ( 5*nn.dssim(gpu_target_src_masked_opt, gpu_pred_src_src_masked_opt, max_val=1.0, filter_size=int(resolution/11.6)), axis=[1]) + gpu_src_loss += tf.reduce_mean ( 5*nn.dssim(gpu_target_src_masked_opt, gpu_pred_src_src_masked_opt, max_val=1.0, filter_size=int(resolution/23.2)), axis=[1]) gpu_src_loss += tf.reduce_mean ( 10*tf.square ( gpu_target_src_masked_opt - gpu_pred_src_src_masked_opt ), axis=[1,2,3]) if eyes_prio: @@ -360,7 +367,11 @@ class SAEHDModel(ModelBase): gpu_src_loss += tf.reduce_mean( (10*bg_style_power)*nn.dssim(gpu_psd_target_dst_anti_masked, gpu_target_dst_anti_masked, max_val=1.0, filter_size=int(resolution/11.6)), axis=[1]) gpu_src_loss += tf.reduce_mean( (10*bg_style_power)*tf.square( gpu_psd_target_dst_anti_masked - gpu_target_dst_anti_masked), axis=[1,2,3] ) - gpu_dst_loss = tf.reduce_mean ( 10*nn.dssim(gpu_target_dst_masked_opt, gpu_pred_dst_dst_masked_opt, max_val=1.0, filter_size=int(resolution/11.6) ), axis=[1]) + if resolution < 256: + gpu_dst_loss = tf.reduce_mean ( 10*nn.dssim(gpu_target_dst_masked_opt, gpu_pred_dst_dst_masked_opt, max_val=1.0, filter_size=int(resolution/11.6) ), axis=[1]) + else: + gpu_dst_loss = tf.reduce_mean ( 5*nn.dssim(gpu_target_dst_masked_opt, gpu_pred_dst_dst_masked_opt, max_val=1.0, filter_size=int(resolution/11.6) ), axis=[1]) + gpu_dst_loss += tf.reduce_mean ( 5*nn.dssim(gpu_target_dst_masked_opt, gpu_pred_dst_dst_masked_opt, max_val=1.0, filter_size=int(resolution/23.2) ), axis=[1]) gpu_dst_loss += tf.reduce_mean ( 10*tf.square( gpu_target_dst_masked_opt- gpu_pred_dst_dst_masked_opt ), axis=[1,2,3]) if eyes_prio: