SAEHD: lr_dropout now as an option, and disabled by default

This commit is contained in:
Colombo 2019-12-20 12:04:44 +04:00
parent 951942821d
commit 021bb6d128

View file

@ -20,7 +20,7 @@ class SAEHDModel(ModelBase):
default_resolution = 128 default_resolution = 128
default_archi = 'df' default_archi = 'df'
default_face_type = 'f' default_face_type = 'f'
if is_first_run: if is_first_run:
resolution = io.input_int("Resolution ( 64-256 ?:help skip:128) : ", default_resolution, help_message="More resolution requires more VRAM and time to train. Value will be adjusted to multiple of 16.") resolution = io.input_int("Resolution ( 64-256 ?:help skip:128) : ", default_resolution, help_message="More resolution requires more VRAM and time to train. Value will be adjusted to multiple of 16.")
@ -38,7 +38,7 @@ class SAEHDModel(ModelBase):
self.options['learn_mask'] = io.input_bool ( f"Learn mask? (y/n, ?:help skip:{yn_str[default_learn_mask]} ) : " , default_learn_mask, help_message="Learning mask can help model to recognize face directions. Learn without mask can reduce model size, in this case converter forced to use 'not predicted mask' that is not smooth as predicted.") self.options['learn_mask'] = io.input_bool ( f"Learn mask? (y/n, ?:help skip:{yn_str[default_learn_mask]} ) : " , default_learn_mask, help_message="Learning mask can help model to recognize face directions. Learn without mask can reduce model size, in this case converter forced to use 'not predicted mask' that is not smooth as predicted.")
else: else:
self.options['learn_mask'] = self.options.get('learn_mask', default_learn_mask) self.options['learn_mask'] = self.options.get('learn_mask', default_learn_mask)
if (is_first_run or ask_override) and 'tensorflow' in self.device_config.backend: if (is_first_run or ask_override) and 'tensorflow' in self.device_config.backend:
def_optimizer_mode = self.options.get('optimizer_mode', 1) def_optimizer_mode = self.options.get('optimizer_mode', 1)
self.options['optimizer_mode'] = io.input_int ("Optimizer mode? ( 1,2,3 ?:help skip:%d) : " % (def_optimizer_mode), def_optimizer_mode, help_message="1 - no changes. 2 - allows you to train x2 bigger network consuming RAM. 3 - allows you to train x3 bigger network consuming huge amount of RAM and slower, depends on CPU power.") self.options['optimizer_mode'] = io.input_int ("Optimizer mode? ( 1,2,3 ?:help skip:%d) : " % (def_optimizer_mode), def_optimizer_mode, help_message="1 - no changes. 2 - allows you to train x2 bigger network consuming RAM. 3 - allows you to train x3 bigger network consuming huge amount of RAM and slower, depends on CPU power.")
@ -63,8 +63,11 @@ class SAEHDModel(ModelBase):
default_true_face_training = self.options.get('true_face_training', False) default_true_face_training = self.options.get('true_face_training', False)
default_face_style_power = self.options.get('face_style_power', 0.0) default_face_style_power = self.options.get('face_style_power', 0.0)
default_bg_style_power = self.options.get('bg_style_power', 0.0) default_bg_style_power = self.options.get('bg_style_power', 0.0)
if is_first_run or ask_override: if is_first_run or ask_override:
default_lr_dropout = self.options.get('lr_dropout', False)
self.options['lr_dropout'] = io.input_bool ( f"Use learning rate dropout? (y/n, ?:help skip:{yn_str[default_lr_dropout]} ) : ", default_lr_dropout, help_message="When the face is trained enough, you can enable this option to get extra sharpness for less amount of iterations.")
default_random_warp = self.options.get('random_warp', True) default_random_warp = self.options.get('random_warp', True)
self.options['random_warp'] = io.input_bool (f"Enable random warp of samples? ( y/n, ?:help skip:{yn_str[default_random_warp]}) : ", 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 for less amount of iterations.") self.options['random_warp'] = io.input_bool (f"Enable random warp of samples? ( y/n, ?:help skip:{yn_str[default_random_warp]}) : ", 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 for less amount of iterations.")
@ -78,14 +81,14 @@ class SAEHDModel(ModelBase):
default_ct_mode = self.options.get('ct_mode', 'none') default_ct_mode = self.options.get('ct_mode', 'none')
self.options['ct_mode'] = io.input_str (f"Color transfer mode apply to src faceset. ( none/rct/lct/mkl/idt/sot, ?:help skip:{default_ct_mode}) : ", default_ct_mode, ['none','rct','lct','mkl','idt','sot'], help_message="Change color distribution of src samples close to dst samples. Try all modes to find the best.") self.options['ct_mode'] = io.input_str (f"Color transfer mode apply to src faceset. ( none/rct/lct/mkl/idt/sot, ?:help skip:{default_ct_mode}) : ", default_ct_mode, ['none','rct','lct','mkl','idt','sot'], help_message="Change color distribution of src samples close to dst samples. Try all modes to find the best.")
if nnlib.device.backend != 'plaidML': # todo https://github.com/plaidml/plaidml/issues/301 if nnlib.device.backend != 'plaidML': # todo https://github.com/plaidml/plaidml/issues/301
default_clipgrad = False if is_first_run else self.options.get('clipgrad', False) default_clipgrad = False if is_first_run else self.options.get('clipgrad', False)
self.options['clipgrad'] = io.input_bool (f"Enable gradient clipping? (y/n, ?:help skip:{yn_str[default_clipgrad]}) : ", default_clipgrad, help_message="Gradient clipping reduces chance of model collapse, sacrificing speed of training.") self.options['clipgrad'] = io.input_bool (f"Enable gradient clipping? (y/n, ?:help skip:{yn_str[default_clipgrad]}) : ", default_clipgrad, help_message="Gradient clipping reduces chance of model collapse, sacrificing speed of training.")
else: else:
self.options['clipgrad'] = False self.options['clipgrad'] = False
else: else:
self.options['lr_dropout'] = self.options.get('lr_dropout', default_lr_dropout)
self.options['random_warp'] = self.options.get('random_warp', True) self.options['random_warp'] = self.options.get('random_warp', True)
self.options['true_face_training'] = self.options.get('true_face_training', default_true_face_training) self.options['true_face_training'] = self.options.get('true_face_training', default_true_face_training)
self.options['face_style_power'] = self.options.get('face_style_power', default_face_style_power) self.options['face_style_power'] = self.options.get('face_style_power', default_face_style_power)
@ -452,10 +455,11 @@ class SAEHDModel(ModelBase):
psd_target_dst_anti_masked = self.model.pred_src_dst*(1.0 - target_dstm) psd_target_dst_anti_masked = self.model.pred_src_dst*(1.0 - target_dstm)
if self.is_training_mode: if self.is_training_mode:
self.src_dst_opt = RMSprop(lr=5e-5, lr_dropout=0.3, clipnorm=1.0 if self.options['clipgrad'] else 0.0, tf_cpu_mode=self.options['optimizer_mode']-1) lr_dropout = 0.3 if self.options['lr_dropout'] else 0.0
self.src_dst_mask_opt = RMSprop(lr=5e-5, lr_dropout=0.3, clipnorm=1.0 if self.options['clipgrad'] else 0.0, tf_cpu_mode=self.options['optimizer_mode']-1) self.src_dst_opt = RMSprop(lr=5e-5, lr_dropout=lr_dropout, clipnorm=1.0 if self.options['clipgrad'] else 0.0, tf_cpu_mode=self.options['optimizer_mode']-1)
self.D_opt = RMSprop(lr=5e-5, lr_dropout=0.3, clipnorm=1.0 if self.options['clipgrad'] else 0.0, tf_cpu_mode=self.options['optimizer_mode']-1) self.src_dst_mask_opt = RMSprop(lr=5e-5, lr_dropout=lr_dropout, clipnorm=1.0 if self.options['clipgrad'] else 0.0, tf_cpu_mode=self.options['optimizer_mode']-1)
self.D_opt = RMSprop(lr=5e-5, lr_dropout=lr_dropout, clipnorm=1.0 if self.options['clipgrad'] else 0.0, tf_cpu_mode=self.options['optimizer_mode']-1)
src_loss = K.mean ( 10*dssim(kernel_size=int(resolution/11.6),max_value=1.0)( target_src_masked_opt, pred_src_src_masked_opt) ) src_loss = K.mean ( 10*dssim(kernel_size=int(resolution/11.6),max_value=1.0)( target_src_masked_opt, pred_src_src_masked_opt) )
src_loss += K.mean ( 10*K.square( target_src_masked_opt - pred_src_src_masked_opt ) ) src_loss += K.mean ( 10*K.square( target_src_masked_opt - pred_src_src_masked_opt ) )
@ -520,7 +524,7 @@ class SAEHDModel(ModelBase):
face_type = t.FACE_TYPE_MID_FULL face_type = t.FACE_TYPE_MID_FULL
elif self.options['face_type'] == 'f': elif self.options['face_type'] == 'f':
face_type = t.FACE_TYPE_FULL face_type = t.FACE_TYPE_FULL
t_mode_bgr = t.MODE_BGR if not self.pretrain else t.MODE_BGR_SHUFFLE t_mode_bgr = t.MODE_BGR if not self.pretrain else t.MODE_BGR_SHUFFLE
training_data_src_path = self.training_data_src_path training_data_src_path = self.training_data_src_path
@ -531,11 +535,11 @@ class SAEHDModel(ModelBase):
training_data_src_path = self.pretraining_data_path training_data_src_path = self.pretraining_data_path
training_data_dst_path = self.pretraining_data_path training_data_dst_path = self.pretraining_data_path
sort_by_yaw = False sort_by_yaw = False
t_img_warped = t.IMG_WARPED_TRANSFORMED if self.options['random_warp'] else t.IMG_TRANSFORMED t_img_warped = t.IMG_WARPED_TRANSFORMED if self.options['random_warp'] else t.IMG_TRANSFORMED
self.set_training_data_generators ([ self.set_training_data_generators ([
SampleGeneratorFace(training_data_src_path, sort_by_yaw_target_samples_path=training_data_dst_path if sort_by_yaw else None, use_caching=False, SampleGeneratorFace(training_data_src_path, sort_by_yaw_target_samples_path=training_data_dst_path if sort_by_yaw else None,
random_ct_samples_path=training_data_dst_path if self.options['ct_mode'] != 'none' else None, random_ct_samples_path=training_data_dst_path if self.options['ct_mode'] != 'none' else None,
debug=self.is_debug(), batch_size=self.batch_size, debug=self.is_debug(), batch_size=self.batch_size,
sample_process_options=SampleProcessor.Options(random_flip=self.random_flip, scale_range=np.array([-0.05, 0.05])+self.src_scale_mod / 100.0 ), sample_process_options=SampleProcessor.Options(random_flip=self.random_flip, scale_range=np.array([-0.05, 0.05])+self.src_scale_mod / 100.0 ),
@ -544,7 +548,7 @@ class SAEHDModel(ModelBase):
{'types' : (t.IMG_TRANSFORMED, face_type, t.MODE_M), 'resolution': resolution } ] {'types' : (t.IMG_TRANSFORMED, face_type, t.MODE_M), 'resolution': resolution } ]
), ),
SampleGeneratorFace(training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size, use_caching=False, SampleGeneratorFace(training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size,
sample_process_options=SampleProcessor.Options(random_flip=self.random_flip, ), sample_process_options=SampleProcessor.Options(random_flip=self.random_flip, ),
output_sample_types = [ {'types' : (t_img_warped, face_type, t_mode_bgr), 'resolution':resolution}, output_sample_types = [ {'types' : (t_img_warped, face_type, t_mode_bgr), 'resolution':resolution},
{'types' : (t.IMG_TRANSFORMED, face_type, t_mode_bgr), 'resolution': resolution}, {'types' : (t.IMG_TRANSFORMED, face_type, t_mode_bgr), 'resolution': resolution},