From f06ff1eae8ded0969803cc2346e00ba9adb4d8bc Mon Sep 17 00:00:00 2001 From: jh Date: Fri, 12 Mar 2021 07:57:09 -0800 Subject: [PATCH 01/44] feat: ms_ssim implementation --- core/leras/ops/__init__.py | 33 +++++++++++++++++++++++++++++++-- 1 file changed, 31 insertions(+), 2 deletions(-) diff --git a/core/leras/ops/__init__.py b/core/leras/ops/__init__.py index 500a22a..af59c6c 100644 --- a/core/leras/ops/__init__.py +++ b/core/leras/ops/__init__.py @@ -307,6 +307,35 @@ def dssim(img1,img2, max_val, filter_size=11, filter_sigma=1.5, k1=0.01, k2=0.03 nn.dssim = dssim + +def ms_ssim(img1, img2, resolution, kernel_size=11, k1=0.01, k2=0.03, max_value=1.0, + power_factors=(0.0448, 0.2856, 0.3001, 0.2363, 0.1333)): + img_dtype = img1.dtype + if img_dtype != img2.dtype: + raise ValueError("img1.dtype != img2.dtype") + + if img_dtype != tf.float32: + img1 = tf.cast(img1, tf.float32) + img2 = tf.cast(img2, tf.float32) + + # restrict mssim factors to those greater/equal to kernel size + power_factors = [power_factors[i] for i in range(len(power_factors)) if resolution//(2**i) >= kernel_size] + + # normalize power factors if reduced because of size + if sum(power_factors) < 1.0: + power_factors = [x/sum(power_factors) for x in power_factors] + + ms_ssim_val = tf.image.ssim_multiscale(img1, img2, max_value=max_value, power_factors=power_factors, + filter_size=kernel_size, k1=k1, k2=k2) + loss = (1.0 - ms_ssim_val) / 2.0 + + if img_dtype != tf.float32: + loss = tf.cast(loss, img_dtype) + return loss + +nn.dssim = ms_ssim + + def space_to_depth(x, size): if nn.data_format == "NHWC": # match NCHW version in order to switch data_format without problems @@ -385,7 +414,7 @@ def total_variation_mse(images): """ pixel_dif1 = images[:, 1:, :, :] - images[:, :-1, :, :] pixel_dif2 = images[:, :, 1:, :] - images[:, :, :-1, :] - + tot_var = ( tf.reduce_sum(tf.square(pixel_dif1), axis=[1,2,3]) + tf.reduce_sum(tf.square(pixel_dif2), axis=[1,2,3]) ) return tot_var @@ -400,4 +429,4 @@ def tf_suppress_lower_mean(t, eps=0.00001): q = tf.clip_by_value(q-t_mean_eps, 0, eps) q = q * (t/eps) return q -""" \ No newline at end of file +""" From 7cb33128f06a00da8f22cadbe4f9a23481779726 Mon Sep 17 00:00:00 2001 From: jh Date: Fri, 12 Mar 2021 08:08:36 -0800 Subject: [PATCH 02/44] feat: add ms_ssim_loss option to SAEHD prompts --- models/Model_SAEHD/Model.py | 28 ++++++++++++++++------------ 1 file changed, 16 insertions(+), 12 deletions(-) diff --git a/models/Model_SAEHD/Model.py b/models/Model_SAEHD/Model.py index 02dbddf..65bca2d 100644 --- a/models/Model_SAEHD/Model.py +++ b/models/Model_SAEHD/Model.py @@ -53,6 +53,8 @@ class SAEHDModel(ModelBase): lr_dropout = {True:'y', False:'n'}.get(lr_dropout, lr_dropout) #backward comp default_lr_dropout = self.options['lr_dropout'] = lr_dropout + default_ms_ssim_loss = self.options['ms_ssim_loss'] = self.load_or_def_option('ms_ssim_loss', False) + default_random_warp = self.options['random_warp'] = self.load_or_def_option('random_warp', True) default_true_face_power = self.options['true_face_power'] = self.load_or_def_option('true_face_power', 0.0) default_face_style_power = self.options['face_style_power'] = self.load_or_def_option('face_style_power', 0.0) @@ -74,7 +76,7 @@ class SAEHDModel(ModelBase): resolution = np.clip ( (resolution // 16) * 16, min_res, max_res) self.options['resolution'] = resolution self.options['face_type'] = io.input_str ("Face type", default_face_type, ['h','mf','f','wf','head', 'custom'], help_message="Half / mid face / full face / whole face / head / custom. Half face has better resolution, but covers less area of cheeks. Mid face is 30% wider than half face. 'Whole face' covers full area of face include forehead. 'head' covers full head, but requires XSeg for src and dst faceset.").lower() - + while True: archi = io.input_str ("AE architecture", default_archi, help_message=\ """ @@ -135,11 +137,11 @@ Examples: df, liae, df-d, df-ud, liae-ud, ... self.options['mouth_prio'] = io.input_bool ("Mouth priority", default_mouth_prio, help_message='Helps to fix mouth problems during training by forcing the neural network to train mouth with higher priority similar to eyes ') self.options['uniform_yaw'] = io.input_bool ("Uniform yaw distribution of samples", default_uniform_yaw, help_message='Helps to fix blurry side faces due to small amount of them in the faceset.') - + default_gan_power = self.options['gan_power'] = self.load_or_def_option('gan_power', 0.0) default_gan_patch_size = self.options['gan_patch_size'] = self.load_or_def_option('gan_patch_size', self.options['resolution'] // 8) default_gan_dims = self.options['gan_dims'] = self.load_or_def_option('gan_dims', 16) - + 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.") @@ -147,17 +149,19 @@ Examples: df, liae, df-d, df-ud, liae-ud, ... 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 less amount of iterations. Enabled it before `disable random warp` and before GAN. \nn - disabled.\ny - enabled\ncpu - enabled on CPU. This allows not to use extra VRAM, sacrificing 20% time of iteration.") + self.options['ms_ssim_loss'] = io.input_bool("Use multiscale loss?", default_ms_ssim_loss, help_message="Use Multiscale structural similarity for image quality assessment.") + 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 less amount of iterations.") self.options['gan_power'] = np.clip ( io.input_number ("GAN power", default_gan_power, add_info="0.0 .. 1.0", help_message="Forces the neural network to learn small details of the face. Enable it only when the face is trained enough with lr_dropout(on) and random_warp(off), and don't disable. The higher the value, the higher the chances of artifacts. Typical fine value is 0.1"), 0.0, 1.0 ) - - if self.options['gan_power'] != 0.0: + + if self.options['gan_power'] != 0.0: gan_patch_size = np.clip ( io.input_int("GAN patch size", default_gan_patch_size, add_info="3-640", help_message="The higher patch size, the higher the quality, the more VRAM is required. You can get sharper edges even at the lowest setting. Typical fine value is resolution / 8." ), 3, 640 ) self.options['gan_patch_size'] = gan_patch_size - + gan_dims = np.clip ( io.input_int("GAN dimensions", default_gan_dims, add_info="4-64", help_message="The dimensions of the GAN network. The higher dimensions, the more VRAM is required. You can get sharper edges even at the lowest setting. Typical fine value is 16." ), 4, 64 ) self.options['gan_dims'] = gan_dims - + if 'df' in self.options['archi']: self.options['true_face_power'] = np.clip ( io.input_number ("'True face' power.", default_true_face_power, add_info="0.0000 .. 1.0", help_message="Experimental option. Discriminates result face to be more like src face. Higher value - stronger discrimination. Typical value is 0.01 . Comparison - https://i.imgur.com/czScS9q.png"), 0.0, 1.0 ) else: @@ -173,7 +177,7 @@ Examples: df, liae, df-d, df-ud, liae-ud, ... if self.options['pretrain'] and self.get_pretraining_data_path() is None: raise Exception("pretraining_data_path is not defined") - + self.gan_model_changed = (default_gan_patch_size != self.options['gan_patch_size']) or (default_gan_dims != self.options['gan_dims']) self.pretrain_just_disabled = (default_pretrain == True and self.options['pretrain'] == False) @@ -406,7 +410,7 @@ Examples: df, liae, df-d, df-ud, liae-ud, ... gpu_target_dstm_style_blur = gpu_target_dstm_blur #default style mask is 0.5 on boundary gpu_target_dstm_blur = tf.clip_by_value(gpu_target_dstm_blur, 0, 0.5) * 2 - gpu_target_dst_masked = gpu_target_dst*gpu_target_dstm_blur + gpu_target_dst_masked = gpu_target_dst*gpu_target_dstm_blur gpu_target_dst_style_masked = gpu_target_dst*gpu_target_dstm_style_blur gpu_target_dst_style_anti_masked = gpu_target_dst*(1.0 - gpu_target_dstm_style_blur) @@ -516,13 +520,13 @@ Examples: df, liae, df-d, df-ud, liae-ud, ... gpu_G_loss += gan_power*(DLoss(gpu_pred_src_src_d_ones, gpu_pred_src_src_d) + \ DLoss(gpu_pred_src_src_d2_ones, gpu_pred_src_src_d2)) - - + + if masked_training: # Minimal src-src-bg rec with total_variation_mse to suppress random bright dots from gan gpu_G_loss += 0.000001*nn.total_variation_mse(gpu_pred_src_src) gpu_G_loss += 0.02*tf.reduce_mean(tf.square(gpu_pred_src_src_anti_masked-gpu_target_src_anti_masked),axis=[1,2,3] ) - + gpu_G_loss_gvs += [ nn.gradients ( gpu_G_loss, self.src_dst_trainable_weights ) ] From 039559e4fdb8cddb1b02f048b53971b6470f0c18 Mon Sep 17 00:00:00 2001 From: jh Date: Fri, 12 Mar 2021 08:11:21 -0800 Subject: [PATCH 03/44] fix: typo in new ms_ssim opt function --- core/leras/ops/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/core/leras/ops/__init__.py b/core/leras/ops/__init__.py index af59c6c..b7f1bab 100644 --- a/core/leras/ops/__init__.py +++ b/core/leras/ops/__init__.py @@ -333,7 +333,7 @@ def ms_ssim(img1, img2, resolution, kernel_size=11, k1=0.01, k2=0.03, max_value= loss = tf.cast(loss, img_dtype) return loss -nn.dssim = ms_ssim +nn.ms_ssim = ms_ssim def space_to_depth(x, size): From d5fd69906d2eac8ba1f5359dfd14a1adf7b8a5b3 Mon Sep 17 00:00:00 2001 From: jh Date: Fri, 12 Mar 2021 08:17:04 -0800 Subject: [PATCH 04/44] feat: first test of ms_ssim --- models/Model_SAEHD/Model.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/models/Model_SAEHD/Model.py b/models/Model_SAEHD/Model.py index 65bca2d..2989d6e 100644 --- a/models/Model_SAEHD/Model.py +++ b/models/Model_SAEHD/Model.py @@ -425,12 +425,15 @@ Examples: df, liae, df-d, df-ud, liae-ud, ... gpu_psd_target_dst_style_masked = gpu_pred_src_dst*gpu_target_dstm_style_blur gpu_psd_target_dst_style_anti_masked = gpu_pred_src_dst*(1.0 - gpu_target_dstm_style_blur) - 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]) + if self.options['ms_ssim_loss']: + gpu_src_loss = tf.reduce_mean ( 10*nn.ms_ssim(gpu_target_src_masked_opt, gpu_pred_src_src_masked_opt, resolution), 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 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 or mouth_prio: if eyes_prio and mouth_prio: From 6b78f1cbc3a0ad4d65bb3c7a428f4ccdab169c18 Mon Sep 17 00:00:00 2001 From: jh Date: Fri, 12 Mar 2021 08:18:56 -0800 Subject: [PATCH 05/44] fix: misnamed arg for tf.image.ssim_multiscale --- core/leras/ops/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/core/leras/ops/__init__.py b/core/leras/ops/__init__.py index b7f1bab..398679c 100644 --- a/core/leras/ops/__init__.py +++ b/core/leras/ops/__init__.py @@ -325,7 +325,7 @@ def ms_ssim(img1, img2, resolution, kernel_size=11, k1=0.01, k2=0.03, max_value= if sum(power_factors) < 1.0: power_factors = [x/sum(power_factors) for x in power_factors] - ms_ssim_val = tf.image.ssim_multiscale(img1, img2, max_value=max_value, power_factors=power_factors, + ms_ssim_val = tf.image.ssim_multiscale(img1, img2, max_val=max_value, power_factors=power_factors, filter_size=kernel_size, k1=k1, k2=k2) loss = (1.0 - ms_ssim_val) / 2.0 From 0df0d6d5cb365238432de2793d92652c82f4a9f6 Mon Sep 17 00:00:00 2001 From: jh Date: Fri, 12 Mar 2021 08:26:14 -0800 Subject: [PATCH 06/44] fix: transpose batches for tf msssim --- core/leras/ops/__init__.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/core/leras/ops/__init__.py b/core/leras/ops/__init__.py index 398679c..6824488 100644 --- a/core/leras/ops/__init__.py +++ b/core/leras/ops/__init__.py @@ -325,7 +325,10 @@ def ms_ssim(img1, img2, resolution, kernel_size=11, k1=0.01, k2=0.03, max_value= if sum(power_factors) < 1.0: power_factors = [x/sum(power_factors) for x in power_factors] - ms_ssim_val = tf.image.ssim_multiscale(img1, img2, max_val=max_value, power_factors=power_factors, + # Transpose images from NCHW to NHWC + img1_t = tf.transpose(img1, [0, 2, 3, 1]) + img2_t = tf.transpose(img2, [0, 2, 3, 1]) + ms_ssim_val = tf.image.ssim_multiscale(img1_t, img2_t, max_val=max_value, power_factors=power_factors, filter_size=kernel_size, k1=k1, k2=k2) loss = (1.0 - ms_ssim_val) / 2.0 From 19461ef76f97a5b507d96c9f6d5c90680ce52d24 Mon Sep 17 00:00:00 2001 From: jh Date: Fri, 12 Mar 2021 08:29:30 -0800 Subject: [PATCH 07/44] fix: remove axis from ms_ssim call --- models/Model_SAEHD/Model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/models/Model_SAEHD/Model.py b/models/Model_SAEHD/Model.py index 2989d6e..2db427d 100644 --- a/models/Model_SAEHD/Model.py +++ b/models/Model_SAEHD/Model.py @@ -426,7 +426,7 @@ Examples: df, liae, df-d, df-ud, liae-ud, ... gpu_psd_target_dst_style_anti_masked = gpu_pred_src_dst*(1.0 - gpu_target_dstm_style_blur) if self.options['ms_ssim_loss']: - gpu_src_loss = tf.reduce_mean ( 10*nn.ms_ssim(gpu_target_src_masked_opt, gpu_pred_src_src_masked_opt, resolution), axis=[1]) + gpu_src_loss = tf.reduce_mean ( 10*nn.ms_ssim(gpu_target_src_masked_opt, gpu_pred_src_src_masked_opt, resolution)) else: 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]) From f5edf30e84d110f7708d92dfab060a70a2557306 Mon Sep 17 00:00:00 2001 From: jh Date: Fri, 12 Mar 2021 08:41:48 -0800 Subject: [PATCH 08/44] fix: try inner function --- core/leras/ops/__init__.py | 38 +++++++++++++++++++++---------------- models/Model_SAEHD/Model.py | 2 +- 2 files changed, 23 insertions(+), 17 deletions(-) diff --git a/core/leras/ops/__init__.py b/core/leras/ops/__init__.py index 6824488..c136009 100644 --- a/core/leras/ops/__init__.py +++ b/core/leras/ops/__init__.py @@ -308,15 +308,8 @@ def dssim(img1,img2, max_val, filter_size=11, filter_sigma=1.5, k1=0.01, k2=0.03 nn.dssim = dssim -def ms_ssim(img1, img2, resolution, kernel_size=11, k1=0.01, k2=0.03, max_value=1.0, +def ms_ssim(resolution, kernel_size=11, k1=0.01, k2=0.03, max_value=1.0, power_factors=(0.0448, 0.2856, 0.3001, 0.2363, 0.1333)): - img_dtype = img1.dtype - if img_dtype != img2.dtype: - raise ValueError("img1.dtype != img2.dtype") - - if img_dtype != tf.float32: - img1 = tf.cast(img1, tf.float32) - img2 = tf.cast(img2, tf.float32) # restrict mssim factors to those greater/equal to kernel size power_factors = [power_factors[i] for i in range(len(power_factors)) if resolution//(2**i) >= kernel_size] @@ -325,15 +318,28 @@ def ms_ssim(img1, img2, resolution, kernel_size=11, k1=0.01, k2=0.03, max_value= if sum(power_factors) < 1.0: power_factors = [x/sum(power_factors) for x in power_factors] - # Transpose images from NCHW to NHWC - img1_t = tf.transpose(img1, [0, 2, 3, 1]) - img2_t = tf.transpose(img2, [0, 2, 3, 1]) - ms_ssim_val = tf.image.ssim_multiscale(img1_t, img2_t, max_val=max_value, power_factors=power_factors, - filter_size=kernel_size, k1=k1, k2=k2) - loss = (1.0 - ms_ssim_val) / 2.0 + def loss(img1, img2): + img_dtype = img1.dtype + if img_dtype != img2.dtype: + raise ValueError("img1.dtype != img2.dtype") + + if img_dtype != tf.float32: + img1 = tf.cast(img1, tf.float32) + img2 = tf.cast(img2, tf.float32) + + + + # Transpose images from NCHW to NHWC + img1_t = tf.transpose(img1, [0, 2, 3, 1]) + img2_t = tf.transpose(img2, [0, 2, 3, 1]) + ms_ssim_val = tf.image.ssim_multiscale(img1_t, img2_t, max_val=max_value, power_factors=power_factors, + filter_size=kernel_size, k1=k1, k2=k2) + ms_ssim_loss = (1.0 - ms_ssim_val) / 2.0 + + if img_dtype != tf.float32: + ms_ssim_loss = tf.cast(ms_ssim_loss, img_dtype) + return ms_ssim_loss - if img_dtype != tf.float32: - loss = tf.cast(loss, img_dtype) return loss nn.ms_ssim = ms_ssim diff --git a/models/Model_SAEHD/Model.py b/models/Model_SAEHD/Model.py index 2db427d..16d63da 100644 --- a/models/Model_SAEHD/Model.py +++ b/models/Model_SAEHD/Model.py @@ -426,7 +426,7 @@ Examples: df, liae, df-d, df-ud, liae-ud, ... gpu_psd_target_dst_style_anti_masked = gpu_pred_src_dst*(1.0 - gpu_target_dstm_style_blur) if self.options['ms_ssim_loss']: - gpu_src_loss = tf.reduce_mean ( 10*nn.ms_ssim(gpu_target_src_masked_opt, gpu_pred_src_src_masked_opt, resolution)) + gpu_src_loss = tf.reduce_mean ( 10*nn.ms_ssim(resolution)(gpu_target_src_masked_opt, gpu_pred_src_src_masked_opt)) else: 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]) From 28279971a17a5300ced4e129b8bf953063007b44 Mon Sep 17 00:00:00 2001 From: jh Date: Fri, 12 Mar 2021 11:28:01 -0800 Subject: [PATCH 09/44] fix: more testing --- core/leras/ops/__init__.py | 37 ++++++++++++++++--------------------- 1 file changed, 16 insertions(+), 21 deletions(-) diff --git a/core/leras/ops/__init__.py b/core/leras/ops/__init__.py index c136009..583a67d 100644 --- a/core/leras/ops/__init__.py +++ b/core/leras/ops/__init__.py @@ -308,7 +308,7 @@ def dssim(img1,img2, max_val, filter_size=11, filter_sigma=1.5, k1=0.01, k2=0.03 nn.dssim = dssim -def ms_ssim(resolution, kernel_size=11, k1=0.01, k2=0.03, max_value=1.0, +def ms_ssim(img1, img2, resolution, kernel_size=11, k1=0.01, k2=0.03, max_value=1.0, power_factors=(0.0448, 0.2856, 0.3001, 0.2363, 0.1333)): # restrict mssim factors to those greater/equal to kernel size @@ -318,29 +318,24 @@ def ms_ssim(resolution, kernel_size=11, k1=0.01, k2=0.03, max_value=1.0, if sum(power_factors) < 1.0: power_factors = [x/sum(power_factors) for x in power_factors] - def loss(img1, img2): - img_dtype = img1.dtype - if img_dtype != img2.dtype: - raise ValueError("img1.dtype != img2.dtype") + img_dtype = img1.dtype + if img_dtype != img2.dtype: + raise ValueError("img1.dtype != img2.dtype") - if img_dtype != tf.float32: - img1 = tf.cast(img1, tf.float32) - img2 = tf.cast(img2, tf.float32) + if img_dtype != tf.float32: + img1 = tf.cast(img1, tf.float32) + img2 = tf.cast(img2, tf.float32) + # Transpose images from NCHW to NHWC + img1_t = tf.transpose(img1, [0, 2, 3, 1]) + img2_t = tf.transpose(img2, [0, 2, 3, 1]) + ms_ssim_val = tf.image.ssim_multiscale(img1_t, img2_t, max_val=max_value, power_factors=power_factors, + filter_size=kernel_size, k1=k1, k2=k2) + ms_ssim_loss = (1.0 - ms_ssim_val) / 2.0 - - # Transpose images from NCHW to NHWC - img1_t = tf.transpose(img1, [0, 2, 3, 1]) - img2_t = tf.transpose(img2, [0, 2, 3, 1]) - ms_ssim_val = tf.image.ssim_multiscale(img1_t, img2_t, max_val=max_value, power_factors=power_factors, - filter_size=kernel_size, k1=k1, k2=k2) - ms_ssim_loss = (1.0 - ms_ssim_val) / 2.0 - - if img_dtype != tf.float32: - ms_ssim_loss = tf.cast(ms_ssim_loss, img_dtype) - return ms_ssim_loss - - return loss + if img_dtype != tf.float32: + ms_ssim_loss = tf.cast(ms_ssim_loss, img_dtype) + return ms_ssim_loss nn.ms_ssim = ms_ssim From 3588b0efa96d5d737a02fb111a5723d7a15f7e5c Mon Sep 17 00:00:00 2001 From: jh Date: Fri, 12 Mar 2021 11:29:44 -0800 Subject: [PATCH 10/44] fix: model --- models/Model_SAEHD/Model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/models/Model_SAEHD/Model.py b/models/Model_SAEHD/Model.py index 00c1fc0..23065ec 100644 --- a/models/Model_SAEHD/Model.py +++ b/models/Model_SAEHD/Model.py @@ -426,7 +426,7 @@ Examples: df, liae, df-d, df-ud, liae-ud, ... gpu_psd_target_dst_style_anti_masked = gpu_pred_src_dst*(1.0 - gpu_target_dstm_style_blur) if self.options['ms_ssim_loss']: - gpu_src_loss = tf.reduce_mean ( 10*nn.ms_ssim(resolution)(gpu_target_src_masked_opt, gpu_pred_src_src_masked_opt)) + gpu_src_loss = tf.reduce_mean ( 10*nn.ms_ssim(gpu_target_src_masked_opt, gpu_pred_src_src_masked_opt, resolution)) else: 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]) From d35549fc5d86ff9e38f56ce06ffc9f4e3652efab Mon Sep 17 00:00:00 2001 From: jh Date: Fri, 12 Mar 2021 16:00:20 -0800 Subject: [PATCH 11/44] Try assign device --- core/leras/ops/__init__.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/core/leras/ops/__init__.py b/core/leras/ops/__init__.py index 583a67d..b460f17 100644 --- a/core/leras/ops/__init__.py +++ b/core/leras/ops/__init__.py @@ -329,8 +329,16 @@ def ms_ssim(img1, img2, resolution, kernel_size=11, k1=0.01, k2=0.03, max_value= # Transpose images from NCHW to NHWC img1_t = tf.transpose(img1, [0, 2, 3, 1]) img2_t = tf.transpose(img2, [0, 2, 3, 1]) - ms_ssim_val = tf.image.ssim_multiscale(img1_t, img2_t, max_val=max_value, power_factors=power_factors, - filter_size=kernel_size, k1=k1, k2=k2) + + def assign_device(op): + if op.type != 'ListDiff': + return '/gpu:0' + else: + return '/cpu:0' + + with tf.device(assign_device): + ms_ssim_val = tf.image.ssim_multiscale(img1_t, img2_t, max_val=max_value, power_factors=power_factors, + filter_size=kernel_size, k1=k1, k2=k2) ms_ssim_loss = (1.0 - ms_ssim_val) / 2.0 if img_dtype != tf.float32: From 9b98607b334136a42ecd7a7173647103dd806331 Mon Sep 17 00:00:00 2001 From: jh Date: Fri, 12 Mar 2021 17:52:57 -0800 Subject: [PATCH 12/44] test: try as layer --- core/leras/layers/MsSsim.py | 25 ++++++++++++ core/leras/ops/__init__.py | 76 ++++++++++++++++++------------------- models/Model_SAEHD/Model.py | 2 +- 3 files changed, 64 insertions(+), 39 deletions(-) create mode 100644 core/leras/layers/MsSsim.py diff --git a/core/leras/layers/MsSsim.py b/core/leras/layers/MsSsim.py new file mode 100644 index 0000000..d1ad64c --- /dev/null +++ b/core/leras/layers/MsSsim.py @@ -0,0 +1,25 @@ +from core.leras import nn +tf = nn.tf + +class MsSsim(nn.LayerBase): + default_power_factors = (0.0448, 0.2856, 0.3001, 0.2363, 0.1333) + + def __init__(self, resolution, kernel_size=11, **kwargs): + # restrict mssim factors to those greater/equal to kernel size + power_factors = [p for i, p in enumerate(self.default_power_factors) if resolution//(2**i) >= kernel_size] + # normalize power factors if reduced because of size + if sum(power_factors) < 1.0: + power_factors = [x/sum(power_factors) for x in power_factors] + self.power_factors = power_factors + self.kernel_size = kernel_size + + super().__init__(**kwargs) + + def __call__(self, y_true, y_pred, max_val): + # Transpose images from NCHW to NHWC + y_true_t = tf.transpose(y_true, [0, 2, 3, 1]) + y_pred_t = tf.transpose(y_pred, [0, 2, 3, 1]) + loss = tf.image.ssim_multiscale(y_true, y_pred, max_val, power_factors=self.power_factors) + return (1.0 - loss) / 2.0 + +nn.MsSsim = MsSsim diff --git a/core/leras/ops/__init__.py b/core/leras/ops/__init__.py index b460f17..f34461f 100644 --- a/core/leras/ops/__init__.py +++ b/core/leras/ops/__init__.py @@ -308,44 +308,44 @@ def dssim(img1,img2, max_val, filter_size=11, filter_sigma=1.5, k1=0.01, k2=0.03 nn.dssim = dssim -def ms_ssim(img1, img2, resolution, kernel_size=11, k1=0.01, k2=0.03, max_value=1.0, - power_factors=(0.0448, 0.2856, 0.3001, 0.2363, 0.1333)): - - # restrict mssim factors to those greater/equal to kernel size - power_factors = [power_factors[i] for i in range(len(power_factors)) if resolution//(2**i) >= kernel_size] - - # normalize power factors if reduced because of size - if sum(power_factors) < 1.0: - power_factors = [x/sum(power_factors) for x in power_factors] - - img_dtype = img1.dtype - if img_dtype != img2.dtype: - raise ValueError("img1.dtype != img2.dtype") - - if img_dtype != tf.float32: - img1 = tf.cast(img1, tf.float32) - img2 = tf.cast(img2, tf.float32) - - # Transpose images from NCHW to NHWC - img1_t = tf.transpose(img1, [0, 2, 3, 1]) - img2_t = tf.transpose(img2, [0, 2, 3, 1]) - - def assign_device(op): - if op.type != 'ListDiff': - return '/gpu:0' - else: - return '/cpu:0' - - with tf.device(assign_device): - ms_ssim_val = tf.image.ssim_multiscale(img1_t, img2_t, max_val=max_value, power_factors=power_factors, - filter_size=kernel_size, k1=k1, k2=k2) - ms_ssim_loss = (1.0 - ms_ssim_val) / 2.0 - - if img_dtype != tf.float32: - ms_ssim_loss = tf.cast(ms_ssim_loss, img_dtype) - return ms_ssim_loss - -nn.ms_ssim = ms_ssim +# def ms_ssim(img1, img2, resolution, kernel_size=11, k1=0.01, k2=0.03, max_value=1.0, +# power_factors=(0.0448, 0.2856, 0.3001, 0.2363, 0.1333)): +# +# # restrict mssim factors to those greater/equal to kernel size +# power_factors = [power_factors[i] for i in range(len(power_factors)) if resolution//(2**i) >= kernel_size] +# +# # normalize power factors if reduced because of size +# if sum(power_factors) < 1.0: +# power_factors = [x/sum(power_factors) for x in power_factors] +# +# img_dtype = img1.dtype +# if img_dtype != img2.dtype: +# raise ValueError("img1.dtype != img2.dtype") +# +# if img_dtype != tf.float32: +# img1 = tf.cast(img1, tf.float32) +# img2 = tf.cast(img2, tf.float32) +# +# # Transpose images from NCHW to NHWC +# img1_t = tf.transpose(img1, [0, 2, 3, 1]) +# img2_t = tf.transpose(img2, [0, 2, 3, 1]) +# +# def assign_device(op): +# if op.type != 'ListDiff': +# return '/gpu:0' +# else: +# return '/cpu:0' +# +# with tf.device(assign_device): +# ms_ssim_val = tf.image.ssim_multiscale(img1_t, img2_t, max_val=max_value, power_factors=power_factors, +# filter_size=kernel_size, k1=k1, k2=k2) +# ms_ssim_loss = (1.0 - ms_ssim_val) / 2.0 +# +# if img_dtype != tf.float32: +# ms_ssim_loss = tf.cast(ms_ssim_loss, img_dtype) +# return ms_ssim_loss +# +# nn.ms_ssim = ms_ssim def space_to_depth(x, size): diff --git a/models/Model_SAEHD/Model.py b/models/Model_SAEHD/Model.py index 23065ec..c57ec6e 100644 --- a/models/Model_SAEHD/Model.py +++ b/models/Model_SAEHD/Model.py @@ -426,7 +426,7 @@ Examples: df, liae, df-d, df-ud, liae-ud, ... gpu_psd_target_dst_style_anti_masked = gpu_pred_src_dst*(1.0 - gpu_target_dstm_style_blur) if self.options['ms_ssim_loss']: - gpu_src_loss = tf.reduce_mean ( 10*nn.ms_ssim(gpu_target_src_masked_opt, gpu_pred_src_src_masked_opt, resolution)) + gpu_src_loss = tf.reduce_mean ( 10*nn.MsSsim(resolution)(gpu_target_src_masked_opt, gpu_pred_src_src_masked_opt)) else: 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]) From 1ac2d26710a5f7d3bca0f2615964146ea449a5db Mon Sep 17 00:00:00 2001 From: jh Date: Fri, 12 Mar 2021 17:58:25 -0800 Subject: [PATCH 13/44] update module --- core/leras/layers/__init__.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/core/leras/layers/__init__.py b/core/leras/layers/__init__.py index 1c81963..d8f1c9d 100644 --- a/core/leras/layers/__init__.py +++ b/core/leras/layers/__init__.py @@ -13,4 +13,5 @@ from .FRNorm2D import * from .TLU import * from .ScaleAdd import * from .DenseNorm import * -from .AdaIN import * \ No newline at end of file +from .AdaIN import * +from .MsSsim import * From 8b8093435f54b51ac34d393df0519fdfe733a186 Mon Sep 17 00:00:00 2001 From: jh Date: Fri, 12 Mar 2021 17:59:55 -0800 Subject: [PATCH 14/44] update args --- models/Model_SAEHD/Model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/models/Model_SAEHD/Model.py b/models/Model_SAEHD/Model.py index c57ec6e..ab1ac1f 100644 --- a/models/Model_SAEHD/Model.py +++ b/models/Model_SAEHD/Model.py @@ -426,7 +426,7 @@ Examples: df, liae, df-d, df-ud, liae-ud, ... gpu_psd_target_dst_style_anti_masked = gpu_pred_src_dst*(1.0 - gpu_target_dstm_style_blur) if self.options['ms_ssim_loss']: - gpu_src_loss = tf.reduce_mean ( 10*nn.MsSsim(resolution)(gpu_target_src_masked_opt, gpu_pred_src_src_masked_opt)) + gpu_src_loss = tf.reduce_mean ( 10*nn.MsSsim(resolution)(gpu_target_src_masked_opt, gpu_pred_src_src_masked_opt, max_val=1.0)) else: 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]) From aaaf9e375d45a4b4e1b64af987fd795dca8354af Mon Sep 17 00:00:00 2001 From: jh Date: Fri, 12 Mar 2021 18:01:34 -0800 Subject: [PATCH 15/44] Fix transposition --- core/leras/layers/MsSsim.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/core/leras/layers/MsSsim.py b/core/leras/layers/MsSsim.py index d1ad64c..12b9ea1 100644 --- a/core/leras/layers/MsSsim.py +++ b/core/leras/layers/MsSsim.py @@ -19,7 +19,7 @@ class MsSsim(nn.LayerBase): # Transpose images from NCHW to NHWC y_true_t = tf.transpose(y_true, [0, 2, 3, 1]) y_pred_t = tf.transpose(y_pred, [0, 2, 3, 1]) - loss = tf.image.ssim_multiscale(y_true, y_pred, max_val, power_factors=self.power_factors) + loss = tf.image.ssim_multiscale(y_true_t, y_pred_t, max_val, power_factors=self.power_factors) return (1.0 - loss) / 2.0 nn.MsSsim = MsSsim From 0225a29832c1a35eacf080e3ea5875a3473bea8a Mon Sep 17 00:00:00 2001 From: jh Date: Fri, 12 Mar 2021 18:07:27 -0800 Subject: [PATCH 16/44] cpu --- core/leras/layers/MsSsim.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/core/leras/layers/MsSsim.py b/core/leras/layers/MsSsim.py index 12b9ea1..d9533f4 100644 --- a/core/leras/layers/MsSsim.py +++ b/core/leras/layers/MsSsim.py @@ -19,7 +19,8 @@ class MsSsim(nn.LayerBase): # Transpose images from NCHW to NHWC y_true_t = tf.transpose(y_true, [0, 2, 3, 1]) y_pred_t = tf.transpose(y_pred, [0, 2, 3, 1]) - loss = tf.image.ssim_multiscale(y_true_t, y_pred_t, max_val, power_factors=self.power_factors) + with tf.device('/CPU:0'): + loss = tf.image.ssim_multiscale(y_true_t, y_pred_t, max_val, power_factors=self.power_factors) return (1.0 - loss) / 2.0 nn.MsSsim = MsSsim From ce5f479f9f2691f1708ecec9b714167fa0e23272 Mon Sep 17 00:00:00 2001 From: jh Date: Fri, 12 Mar 2021 18:09:51 -0800 Subject: [PATCH 17/44] try casting --- core/leras/layers/MsSsim.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/core/leras/layers/MsSsim.py b/core/leras/layers/MsSsim.py index d9533f4..6fa7c69 100644 --- a/core/leras/layers/MsSsim.py +++ b/core/leras/layers/MsSsim.py @@ -17,10 +17,9 @@ class MsSsim(nn.LayerBase): def __call__(self, y_true, y_pred, max_val): # Transpose images from NCHW to NHWC - y_true_t = tf.transpose(y_true, [0, 2, 3, 1]) - y_pred_t = tf.transpose(y_pred, [0, 2, 3, 1]) - with tf.device('/CPU:0'): - loss = tf.image.ssim_multiscale(y_true_t, y_pred_t, max_val, power_factors=self.power_factors) + y_true_t = tf.transpose(tf.cast(y_true), [0, 2, 3, 1]) + y_pred_t = tf.transpose(tf.cast(y_pred), [0, 2, 3, 1]) + loss = tf.image.ssim_multiscale(y_true_t, y_pred_t, max_val, power_factors=self.power_factors) return (1.0 - loss) / 2.0 nn.MsSsim = MsSsim From 6c42cd3a8f21ab7052576ef8f6c932a07101e269 Mon Sep 17 00:00:00 2001 From: jh Date: Fri, 12 Mar 2021 18:11:05 -0800 Subject: [PATCH 18/44] fix cast --- core/leras/layers/MsSsim.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/core/leras/layers/MsSsim.py b/core/leras/layers/MsSsim.py index 6fa7c69..fff5f74 100644 --- a/core/leras/layers/MsSsim.py +++ b/core/leras/layers/MsSsim.py @@ -17,8 +17,8 @@ class MsSsim(nn.LayerBase): def __call__(self, y_true, y_pred, max_val): # Transpose images from NCHW to NHWC - y_true_t = tf.transpose(tf.cast(y_true), [0, 2, 3, 1]) - y_pred_t = tf.transpose(tf.cast(y_pred), [0, 2, 3, 1]) + y_true_t = tf.transpose(tf.cast(y_true, tf.float32), [0, 2, 3, 1]) + y_pred_t = tf.transpose(tf.cast(y_pred, tf.float32), [0, 2, 3, 1]) loss = tf.image.ssim_multiscale(y_true_t, y_pred_t, max_val, power_factors=self.power_factors) return (1.0 - loss) / 2.0 From beb8c30970b310ce1fac421e1448726451a80845 Mon Sep 17 00:00:00 2001 From: jh Date: Fri, 12 Mar 2021 18:27:00 -0800 Subject: [PATCH 19/44] assign device --- core/leras/layers/MsSsim.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/core/leras/layers/MsSsim.py b/core/leras/layers/MsSsim.py index fff5f74..84028bb 100644 --- a/core/leras/layers/MsSsim.py +++ b/core/leras/layers/MsSsim.py @@ -19,7 +19,16 @@ class MsSsim(nn.LayerBase): # Transpose images from NCHW to NHWC y_true_t = tf.transpose(tf.cast(y_true, tf.float32), [0, 2, 3, 1]) y_pred_t = tf.transpose(tf.cast(y_pred, tf.float32), [0, 2, 3, 1]) - loss = tf.image.ssim_multiscale(y_true_t, y_pred_t, max_val, power_factors=self.power_factors) - return (1.0 - loss) / 2.0 + + + def assign_device(op): + if op.type != 'Assert': + return '/gpu:0' + else: + return '/cpu:0' + + with tf.device(assign_device): + loss = tf.image.ssim_multiscale(y_true_t, y_pred_t, max_val, power_factors=self.power_factors) + return (1.0 - loss) / 2.0 nn.MsSsim = MsSsim From 39601603ef9a21abf5fbc571726eeca623307b16 Mon Sep 17 00:00:00 2001 From: jh Date: Fri, 12 Mar 2021 18:32:49 -0800 Subject: [PATCH 20/44] no message --- core/leras/layers/MsSsim.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/core/leras/layers/MsSsim.py b/core/leras/layers/MsSsim.py index 84028bb..b631491 100644 --- a/core/leras/layers/MsSsim.py +++ b/core/leras/layers/MsSsim.py @@ -22,7 +22,7 @@ class MsSsim(nn.LayerBase): def assign_device(op): - if op.type != 'Assert': + if op.type != 'Assert' or op.type != 'ListDiff': return '/gpu:0' else: return '/cpu:0' From 6a68dd5e5954f0ba394494d5dbf913b73b45b6d3 Mon Sep 17 00:00:00 2001 From: jh Date: Fri, 12 Mar 2021 18:36:41 -0800 Subject: [PATCH 21/44] tf2 --- core/leras/layers/MsSsim.py | 13 +++---------- core/leras/nn.py | 19 +++++++++++-------- 2 files changed, 14 insertions(+), 18 deletions(-) diff --git a/core/leras/layers/MsSsim.py b/core/leras/layers/MsSsim.py index b631491..1f7ad6d 100644 --- a/core/leras/layers/MsSsim.py +++ b/core/leras/layers/MsSsim.py @@ -1,5 +1,6 @@ from core.leras import nn tf = nn.tf +tf2 = nn.tf2 class MsSsim(nn.LayerBase): default_power_factors = (0.0448, 0.2856, 0.3001, 0.2363, 0.1333) @@ -20,15 +21,7 @@ class MsSsim(nn.LayerBase): y_true_t = tf.transpose(tf.cast(y_true, tf.float32), [0, 2, 3, 1]) y_pred_t = tf.transpose(tf.cast(y_pred, tf.float32), [0, 2, 3, 1]) - - def assign_device(op): - if op.type != 'Assert' or op.type != 'ListDiff': - return '/gpu:0' - else: - return '/cpu:0' - - with tf.device(assign_device): - loss = tf.image.ssim_multiscale(y_true_t, y_pred_t, max_val, power_factors=self.power_factors) - return (1.0 - loss) / 2.0 + loss = tf2.image.ssim_multiscale(y_true_t, y_pred_t, max_val, power_factors=self.power_factors) + return (1.0 - loss) / 2.0 nn.MsSsim = MsSsim diff --git a/core/leras/nn.py b/core/leras/nn.py index ef5c2c9..8b10fba 100644 --- a/core/leras/nn.py +++ b/core/leras/nn.py @@ -31,6 +31,7 @@ class nn(): current_DeviceConfig = None tf = None + tf2 = None tf_sess = None tf_sess_config = None tf_default_device = None @@ -40,7 +41,7 @@ class nn(): conv2d_spatial_axes = None floatx = None - + @staticmethod def initialize(device_config=None, floatx="float32", data_format="NHWC"): @@ -50,7 +51,7 @@ class nn(): nn.setCurrentDeviceConfig(device_config) # Manipulate environment variables before import tensorflow - + if 'CUDA_VISIBLE_DEVICES' in os.environ.keys(): os.environ.pop('CUDA_VISIBLE_DEVICES') @@ -77,15 +78,17 @@ class nn(): io.log_info("Caching GPU kernels...") import tensorflow - + tf_version = getattr(tensorflow,'VERSION', None) if tf_version is None: tf_version = tensorflow.version.GIT_VERSION if tf_version[0] == 'v': tf_version = tf_version[1:] - + if tf_version[0] == '2': tf = tensorflow.compat.v1 + tf2 = tensorflow + nn.tf2 = tf2 else: tf = tensorflow @@ -93,7 +96,7 @@ class nn(): # Disable tensorflow warnings tf_logger = logging.getLogger('tensorflow') tf_logger.setLevel(logging.ERROR) - + if tf_version[0] == '2': tf.disable_v2_behavior() nn.tf = tf @@ -105,7 +108,7 @@ class nn(): import core.leras.optimizers import core.leras.models import core.leras.archis - + # Configure tensorflow session-config if len(device_config.devices) == 0: nn.tf_default_device = "/CPU:0" @@ -118,7 +121,7 @@ class nn(): config.gpu_options.force_gpu_compatible = True config.gpu_options.allow_growth = True nn.tf_sess_config = config - + if nn.tf_sess is None: nn.tf_sess = tf.Session(config=nn.tf_sess_config) @@ -273,7 +276,7 @@ class nn(): @staticmethod def ask_choose_device(*args, **kwargs): return nn.DeviceConfig.GPUIndexes( nn.ask_choose_device_idxs(*args,**kwargs) ) - + def __init__ (self, devices=None): devices = devices or [] From 2e4fdf20ca545161be8553b0484817c38c13dc65 Mon Sep 17 00:00:00 2001 From: jh Date: Fri, 12 Mar 2021 18:45:34 -0800 Subject: [PATCH 22/44] soft placement --- core/leras/layers/MsSsim.py | 3 +-- core/leras/nn.py | 5 +---- 2 files changed, 2 insertions(+), 6 deletions(-) diff --git a/core/leras/layers/MsSsim.py b/core/leras/layers/MsSsim.py index 1f7ad6d..b8d4c45 100644 --- a/core/leras/layers/MsSsim.py +++ b/core/leras/layers/MsSsim.py @@ -1,6 +1,5 @@ from core.leras import nn tf = nn.tf -tf2 = nn.tf2 class MsSsim(nn.LayerBase): default_power_factors = (0.0448, 0.2856, 0.3001, 0.2363, 0.1333) @@ -21,7 +20,7 @@ class MsSsim(nn.LayerBase): y_true_t = tf.transpose(tf.cast(y_true, tf.float32), [0, 2, 3, 1]) y_pred_t = tf.transpose(tf.cast(y_pred, tf.float32), [0, 2, 3, 1]) - loss = tf2.image.ssim_multiscale(y_true_t, y_pred_t, max_val, power_factors=self.power_factors) + loss = tf.image.ssim_multiscale(y_true_t, y_pred_t, max_val, power_factors=self.power_factors) return (1.0 - loss) / 2.0 nn.MsSsim = MsSsim diff --git a/core/leras/nn.py b/core/leras/nn.py index 8b10fba..8ac437b 100644 --- a/core/leras/nn.py +++ b/core/leras/nn.py @@ -31,7 +31,6 @@ class nn(): current_DeviceConfig = None tf = None - tf2 = None tf_sess = None tf_sess_config = None tf_default_device = None @@ -87,8 +86,6 @@ class nn(): if tf_version[0] == '2': tf = tensorflow.compat.v1 - tf2 = tensorflow - nn.tf2 = tf2 else: tf = tensorflow @@ -115,7 +112,7 @@ class nn(): config = tf.ConfigProto(device_count={'GPU': 0}) else: nn.tf_default_device = "/GPU:0" - config = tf.ConfigProto() + config = tf.ConfigProto(allow_soft_placement=True) config.gpu_options.visible_device_list = ','.join([str(device.index) for device in device_config.devices]) config.gpu_options.force_gpu_compatible = True From 9537c82213cc0a5987e46e32a07be8fdecaf53f3 Mon Sep 17 00:00:00 2001 From: jh Date: Fri, 12 Mar 2021 18:57:24 -0800 Subject: [PATCH 23/44] ms ssim --- models/Model_SAEHD/Model.py | 33 ++++++++++++++++++++++++--------- 1 file changed, 24 insertions(+), 9 deletions(-) diff --git a/models/Model_SAEHD/Model.py b/models/Model_SAEHD/Model.py index ab1ac1f..d97a408 100644 --- a/models/Model_SAEHD/Model.py +++ b/models/Model_SAEHD/Model.py @@ -443,9 +443,15 @@ Examples: df, liae, df-d, df-ud, liae-ud, ... elif mouth_prio: gpu_target_part_mask = gpu_target_srcm_mouth - gpu_src_loss += tf.reduce_mean ( 300*tf.abs ( gpu_target_src*gpu_target_part_mask - gpu_pred_src_src*gpu_target_part_mask ), axis=[1,2,3]) + if self.options['ms_ssim_loss']: + gpu_src_loss += tf.reduce_mean ( 10*nn.MsSsim(resolution)(gpu_target_src*gpu_target_part_mask, gpu_pred_src_src*gpu_target_part_mask, max_val=1.0)) + else: + gpu_src_loss += tf.reduce_mean ( 300*tf.abs ( gpu_target_src*gpu_target_part_mask - gpu_pred_src_src*gpu_target_part_mask ), axis=[1,2,3]) - gpu_src_loss += tf.reduce_mean ( 10*tf.square( gpu_target_srcm - gpu_pred_src_srcm ),axis=[1,2,3] ) + if self.options['ms_ssim_loss']: + gpu_src_loss += tf.reduce_mean ( 10*nn.MsSsim(resolution)(gpu_target_srcm, gpu_pred_src_srcm, max_val=1.0)) + else: + gpu_src_loss += tf.reduce_mean ( 10*tf.square( gpu_target_srcm - gpu_pred_src_srcm ),axis=[1,2,3] ) face_style_power = self.options['face_style_power'] / 100.0 if face_style_power != 0 and not self.pretrain: @@ -456,12 +462,15 @@ Examples: df, liae, df-d, df-ud, liae-ud, ... gpu_src_loss += tf.reduce_mean( (10*bg_style_power)*nn.dssim( gpu_psd_target_dst_style_anti_masked, gpu_target_dst_style_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_style_anti_masked - gpu_target_dst_style_anti_masked), axis=[1,2,3] ) - 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]) + if self.options['ms_ssim_loss']: + gpu_dst_loss = tf.reduce_mean ( 10*nn.MsSsim(resolution)(gpu_target_dst_masked_opt, gpu_pred_dst_dst_masked_opt, max_val=1.0)) 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 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 or mouth_prio: @@ -472,9 +481,15 @@ Examples: df, liae, df-d, df-ud, liae-ud, ... elif mouth_prio: gpu_target_part_mask = gpu_target_dstm_mouth - gpu_dst_loss += tf.reduce_mean ( 300*tf.abs ( gpu_target_dst*gpu_target_part_mask - gpu_pred_dst_dst*gpu_target_part_mask ), axis=[1,2,3]) + if self.options['ms_ssim_loss']: + gpu_dst_loss += tf.reduce_mean ( 10*nn.MsSsim(resolution)(gpu_target_dst*gpu_target_part_mask, gpu_pred_dst_dst*gpu_target_part_mask, max_val=1.0)) + else: + gpu_dst_loss += tf.reduce_mean ( 300*tf.abs ( gpu_target_dst*gpu_target_part_mask - gpu_pred_dst_dst*gpu_target_part_mask ), axis=[1,2,3]) - gpu_dst_loss += tf.reduce_mean ( 10*tf.square( gpu_target_dstm - gpu_pred_dst_dstm ),axis=[1,2,3] ) + if self.options['ms_ssim_loss']: + gpu_dst_loss += tf.reduce_mean ( 10*nn.MsSsim(resolution)(gpu_target_dst_masked_opt, gpu_pred_dst_dst_masked_opt, max_val=1.0)) + else: + gpu_dst_loss += tf.reduce_mean ( 10*tf.square( gpu_target_dstm - gpu_pred_dst_dstm ),axis=[1,2,3] ) gpu_src_losses += [gpu_src_loss] gpu_dst_losses += [gpu_dst_loss] From cddfa763f1371f74eec06aa00c09df1a57fb3464 Mon Sep 17 00:00:00 2001 From: jh Date: Fri, 12 Mar 2021 19:09:14 -0800 Subject: [PATCH 24/44] fix --- models/Model_SAEHD/Model.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/models/Model_SAEHD/Model.py b/models/Model_SAEHD/Model.py index d97a408..13c7e0a 100644 --- a/models/Model_SAEHD/Model.py +++ b/models/Model_SAEHD/Model.py @@ -482,12 +482,12 @@ Examples: df, liae, df-d, df-ud, liae-ud, ... gpu_target_part_mask = gpu_target_dstm_mouth if self.options['ms_ssim_loss']: - gpu_dst_loss += tf.reduce_mean ( 10*nn.MsSsim(resolution)(gpu_target_dst*gpu_target_part_mask, gpu_pred_dst_dst*gpu_target_part_mask, max_val=1.0)) + gpu_dst_loss += tf.reduce_mean ( 10*nn.MsSsim(resolution)(gpu_target_dst*gpu_target_part_mask, gpu_pred_dst_dst*gpu_target_part_mask, max_val=1.0)) else: gpu_dst_loss += tf.reduce_mean ( 300*tf.abs ( gpu_target_dst*gpu_target_part_mask - gpu_pred_dst_dst*gpu_target_part_mask ), axis=[1,2,3]) if self.options['ms_ssim_loss']: - gpu_dst_loss += tf.reduce_mean ( 10*nn.MsSsim(resolution)(gpu_target_dst_masked_opt, gpu_pred_dst_dst_masked_opt, max_val=1.0)) + gpu_dst_loss += tf.reduce_mean ( 10*nn.MsSsim(resolution)(gpu_target_dstm, gpu_pred_dst_dstm, max_val=1.0)) else: gpu_dst_loss += tf.reduce_mean ( 10*tf.square( gpu_target_dstm - gpu_pred_dst_dstm ),axis=[1,2,3] ) From 023be4cf38d1cdc51ddf7ce43732c24f1b964d63 Mon Sep 17 00:00:00 2001 From: jh Date: Fri, 12 Mar 2021 19:13:06 -0800 Subject: [PATCH 25/44] commented out code --- models/Model_SAEHD/Model.py | 32 ++++++++++++++++---------------- 1 file changed, 16 insertions(+), 16 deletions(-) diff --git a/models/Model_SAEHD/Model.py b/models/Model_SAEHD/Model.py index 13c7e0a..c6aa24e 100644 --- a/models/Model_SAEHD/Model.py +++ b/models/Model_SAEHD/Model.py @@ -443,15 +443,15 @@ Examples: df, liae, df-d, df-ud, liae-ud, ... elif mouth_prio: gpu_target_part_mask = gpu_target_srcm_mouth - if self.options['ms_ssim_loss']: - gpu_src_loss += tf.reduce_mean ( 10*nn.MsSsim(resolution)(gpu_target_src*gpu_target_part_mask, gpu_pred_src_src*gpu_target_part_mask, max_val=1.0)) - else: - gpu_src_loss += tf.reduce_mean ( 300*tf.abs ( gpu_target_src*gpu_target_part_mask - gpu_pred_src_src*gpu_target_part_mask ), axis=[1,2,3]) + # if self.options['ms_ssim_loss']: + # gpu_src_loss += tf.reduce_mean ( 10*nn.MsSsim(resolution)(gpu_target_src*gpu_target_part_mask, gpu_pred_src_src*gpu_target_part_mask, max_val=1.0)) + # else: + gpu_src_loss += tf.reduce_mean ( 300*tf.abs ( gpu_target_src*gpu_target_part_mask - gpu_pred_src_src*gpu_target_part_mask ), axis=[1,2,3]) - if self.options['ms_ssim_loss']: - gpu_src_loss += tf.reduce_mean ( 10*nn.MsSsim(resolution)(gpu_target_srcm, gpu_pred_src_srcm, max_val=1.0)) - else: - gpu_src_loss += tf.reduce_mean ( 10*tf.square( gpu_target_srcm - gpu_pred_src_srcm ),axis=[1,2,3] ) + # if self.options['ms_ssim_loss']: + # gpu_src_loss += tf.reduce_mean ( 10*nn.MsSsim(resolution)(gpu_target_srcm, gpu_pred_src_srcm, max_val=1.0)) + # else: + gpu_src_loss += tf.reduce_mean ( 10*tf.square( gpu_target_srcm - gpu_pred_src_srcm ),axis=[1,2,3] ) face_style_power = self.options['face_style_power'] / 100.0 if face_style_power != 0 and not self.pretrain: @@ -481,15 +481,15 @@ Examples: df, liae, df-d, df-ud, liae-ud, ... elif mouth_prio: gpu_target_part_mask = gpu_target_dstm_mouth - if self.options['ms_ssim_loss']: - gpu_dst_loss += tf.reduce_mean ( 10*nn.MsSsim(resolution)(gpu_target_dst*gpu_target_part_mask, gpu_pred_dst_dst*gpu_target_part_mask, max_val=1.0)) - else: - gpu_dst_loss += tf.reduce_mean ( 300*tf.abs ( gpu_target_dst*gpu_target_part_mask - gpu_pred_dst_dst*gpu_target_part_mask ), axis=[1,2,3]) + # if self.options['ms_ssim_loss']: + # gpu_dst_loss += tf.reduce_mean ( 10*nn.MsSsim(resolution)(gpu_target_dst*gpu_target_part_mask, gpu_pred_dst_dst*gpu_target_part_mask, max_val=1.0)) + # else: + gpu_dst_loss += tf.reduce_mean ( 300*tf.abs ( gpu_target_dst*gpu_target_part_mask - gpu_pred_dst_dst*gpu_target_part_mask ), axis=[1,2,3]) - if self.options['ms_ssim_loss']: - gpu_dst_loss += tf.reduce_mean ( 10*nn.MsSsim(resolution)(gpu_target_dstm, gpu_pred_dst_dstm, max_val=1.0)) - else: - gpu_dst_loss += tf.reduce_mean ( 10*tf.square( gpu_target_dstm - gpu_pred_dst_dstm ),axis=[1,2,3] ) + # if self.options['ms_ssim_loss']: + # gpu_dst_loss += tf.reduce_mean ( 10*nn.MsSsim(resolution)(gpu_target_dstm, gpu_pred_dst_dstm, max_val=1.0)) + # else: + gpu_dst_loss += tf.reduce_mean ( 10*tf.square( gpu_target_dstm - gpu_pred_dst_dstm ),axis=[1,2,3] ) gpu_src_losses += [gpu_src_loss] gpu_dst_losses += [gpu_dst_loss] From 1785ecc77d88859ca9451bfa8c4ba7a4a20cd452 Mon Sep 17 00:00:00 2001 From: jh Date: Sat, 13 Mar 2021 12:54:44 -0800 Subject: [PATCH 26/44] ms-ssim loss for eyes/mouth --- models/Model_SAEHD/Model.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/models/Model_SAEHD/Model.py b/models/Model_SAEHD/Model.py index c6aa24e..1d88fbe 100644 --- a/models/Model_SAEHD/Model.py +++ b/models/Model_SAEHD/Model.py @@ -443,10 +443,10 @@ Examples: df, liae, df-d, df-ud, liae-ud, ... elif mouth_prio: gpu_target_part_mask = gpu_target_srcm_mouth - # if self.options['ms_ssim_loss']: - # gpu_src_loss += tf.reduce_mean ( 10*nn.MsSsim(resolution)(gpu_target_src*gpu_target_part_mask, gpu_pred_src_src*gpu_target_part_mask, max_val=1.0)) - # else: - gpu_src_loss += tf.reduce_mean ( 300*tf.abs ( gpu_target_src*gpu_target_part_mask - gpu_pred_src_src*gpu_target_part_mask ), axis=[1,2,3]) + if self.options['ms_ssim_loss']: + gpu_src_loss += tf.reduce_mean ( 300*nn.MsSsim(resolution, kernel_size=5)(gpu_target_src*gpu_target_part_mask, gpu_pred_src_src*gpu_target_part_mask, max_val=1.0)) + else: + gpu_src_loss += tf.reduce_mean ( 300*tf.abs ( gpu_target_src*gpu_target_part_mask - gpu_pred_src_src*gpu_target_part_mask ), axis=[1,2,3]) # if self.options['ms_ssim_loss']: # gpu_src_loss += tf.reduce_mean ( 10*nn.MsSsim(resolution)(gpu_target_srcm, gpu_pred_src_srcm, max_val=1.0)) @@ -481,10 +481,10 @@ Examples: df, liae, df-d, df-ud, liae-ud, ... elif mouth_prio: gpu_target_part_mask = gpu_target_dstm_mouth - # if self.options['ms_ssim_loss']: - # gpu_dst_loss += tf.reduce_mean ( 10*nn.MsSsim(resolution)(gpu_target_dst*gpu_target_part_mask, gpu_pred_dst_dst*gpu_target_part_mask, max_val=1.0)) - # else: - gpu_dst_loss += tf.reduce_mean ( 300*tf.abs ( gpu_target_dst*gpu_target_part_mask - gpu_pred_dst_dst*gpu_target_part_mask ), axis=[1,2,3]) + if self.options['ms_ssim_loss']: + gpu_dst_loss += tf.reduce_mean ( 300*nn.MsSsim(resolution, kernel_size=5)(gpu_target_dst*gpu_target_part_mask, gpu_pred_dst_dst*gpu_target_part_mask, max_val=1.0)) + else: + gpu_dst_loss += tf.reduce_mean ( 300*tf.abs ( gpu_target_dst*gpu_target_part_mask - gpu_pred_dst_dst*gpu_target_part_mask ), axis=[1,2,3]) # if self.options['ms_ssim_loss']: # gpu_dst_loss += tf.reduce_mean ( 10*nn.MsSsim(resolution)(gpu_target_dstm, gpu_pred_dst_dstm, max_val=1.0)) From 703c790e67d110e047a1750fd2d6f34cb5085fc1 Mon Sep 17 00:00:00 2001 From: jh Date: Sat, 13 Mar 2021 14:39:36 -0800 Subject: [PATCH 27/44] kernel size fix --- core/leras/layers/MsSsim.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/core/leras/layers/MsSsim.py b/core/leras/layers/MsSsim.py index b8d4c45..9aebf8f 100644 --- a/core/leras/layers/MsSsim.py +++ b/core/leras/layers/MsSsim.py @@ -20,7 +20,7 @@ class MsSsim(nn.LayerBase): y_true_t = tf.transpose(tf.cast(y_true, tf.float32), [0, 2, 3, 1]) y_pred_t = tf.transpose(tf.cast(y_pred, tf.float32), [0, 2, 3, 1]) - loss = tf.image.ssim_multiscale(y_true_t, y_pred_t, max_val, power_factors=self.power_factors) + loss = tf.image.ssim_multiscale(y_true_t, y_pred_t, max_val, power_factors=self.power_factors, filter_size=self.kernel_size) return (1.0 - loss) / 2.0 nn.MsSsim = MsSsim From fd3b4ad8933ad3e47f2ac998d8c1a2050ad665cb Mon Sep 17 00:00:00 2001 From: jh Date: Sat, 13 Mar 2021 17:25:18 -0800 Subject: [PATCH 28/44] uncomment lines --- models/Model_SAEHD/Model.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/models/Model_SAEHD/Model.py b/models/Model_SAEHD/Model.py index 1d88fbe..afdba77 100644 --- a/models/Model_SAEHD/Model.py +++ b/models/Model_SAEHD/Model.py @@ -448,10 +448,10 @@ Examples: df, liae, df-d, df-ud, liae-ud, ... else: gpu_src_loss += tf.reduce_mean ( 300*tf.abs ( gpu_target_src*gpu_target_part_mask - gpu_pred_src_src*gpu_target_part_mask ), axis=[1,2,3]) - # if self.options['ms_ssim_loss']: - # gpu_src_loss += tf.reduce_mean ( 10*nn.MsSsim(resolution)(gpu_target_srcm, gpu_pred_src_srcm, max_val=1.0)) - # else: - gpu_src_loss += tf.reduce_mean ( 10*tf.square( gpu_target_srcm - gpu_pred_src_srcm ),axis=[1,2,3] ) + if self.options['ms_ssim_loss']: + gpu_src_loss += tf.reduce_mean ( 10*nn.MsSsim(resolution)(gpu_target_srcm, gpu_pred_src_srcm, max_val=1.0)) + else: + gpu_src_loss += tf.reduce_mean ( 10*tf.square( gpu_target_srcm - gpu_pred_src_srcm ),axis=[1,2,3] ) face_style_power = self.options['face_style_power'] / 100.0 if face_style_power != 0 and not self.pretrain: @@ -486,10 +486,10 @@ Examples: df, liae, df-d, df-ud, liae-ud, ... else: gpu_dst_loss += tf.reduce_mean ( 300*tf.abs ( gpu_target_dst*gpu_target_part_mask - gpu_pred_dst_dst*gpu_target_part_mask ), axis=[1,2,3]) - # if self.options['ms_ssim_loss']: - # gpu_dst_loss += tf.reduce_mean ( 10*nn.MsSsim(resolution)(gpu_target_dstm, gpu_pred_dst_dstm, max_val=1.0)) - # else: - gpu_dst_loss += tf.reduce_mean ( 10*tf.square( gpu_target_dstm - gpu_pred_dst_dstm ),axis=[1,2,3] ) + if self.options['ms_ssim_loss']: + gpu_dst_loss += tf.reduce_mean ( 10*nn.MsSsim(resolution)(gpu_target_dstm, gpu_pred_dst_dstm, max_val=1.0)) + else: + gpu_dst_loss += tf.reduce_mean ( 10*tf.square( gpu_target_dstm - gpu_pred_dst_dstm ),axis=[1,2,3] ) gpu_src_losses += [gpu_src_loss] gpu_dst_losses += [gpu_dst_loss] From 33ea98012c0973a717df3b7df59d3091dc331539 Mon Sep 17 00:00:00 2001 From: jh Date: Wed, 17 Mar 2021 11:15:33 -0700 Subject: [PATCH 29/44] fix loss --- core/leras/layers/MsSsim.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/core/leras/layers/MsSsim.py b/core/leras/layers/MsSsim.py index 9aebf8f..6c8e931 100644 --- a/core/leras/layers/MsSsim.py +++ b/core/leras/layers/MsSsim.py @@ -1,6 +1,7 @@ from core.leras import nn tf = nn.tf + class MsSsim(nn.LayerBase): default_power_factors = (0.0448, 0.2856, 0.3001, 0.2363, 0.1333) @@ -20,7 +21,9 @@ class MsSsim(nn.LayerBase): y_true_t = tf.transpose(tf.cast(y_true, tf.float32), [0, 2, 3, 1]) y_pred_t = tf.transpose(tf.cast(y_pred, tf.float32), [0, 2, 3, 1]) - loss = tf.image.ssim_multiscale(y_true_t, y_pred_t, max_val, power_factors=self.power_factors, filter_size=self.kernel_size) - return (1.0 - loss) / 2.0 + ms_ssim_val = tf.image.ssim_multiscale(y_true_t, y_pred_t, max_val, power_factors=self.power_factors, filter_size=self.kernel_size) + # ssim_multiscale returns values in range [0, 1] (where 1 is completely identical) + # subtract from 1 to get loss + return 1 - ms_ssim_val nn.MsSsim = MsSsim From 607863f3fab644a275c11886f74cd32d7048bc6c Mon Sep 17 00:00:00 2001 From: jh Date: Wed, 17 Mar 2021 11:22:02 -0700 Subject: [PATCH 30/44] fix? --- core/leras/layers/MsSsim.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/core/leras/layers/MsSsim.py b/core/leras/layers/MsSsim.py index 6c8e931..cb48294 100644 --- a/core/leras/layers/MsSsim.py +++ b/core/leras/layers/MsSsim.py @@ -24,6 +24,6 @@ class MsSsim(nn.LayerBase): ms_ssim_val = tf.image.ssim_multiscale(y_true_t, y_pred_t, max_val, power_factors=self.power_factors, filter_size=self.kernel_size) # ssim_multiscale returns values in range [0, 1] (where 1 is completely identical) # subtract from 1 to get loss - return 1 - ms_ssim_val + return 1.0 - ms_ssim_val nn.MsSsim = MsSsim From 5f10e3c67e4296066d78263c34a81ba95b894202 Mon Sep 17 00:00:00 2001 From: jh Date: Wed, 17 Mar 2021 11:32:05 -0700 Subject: [PATCH 31/44] revert changes --- core/leras/layers/MsSsim.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/core/leras/layers/MsSsim.py b/core/leras/layers/MsSsim.py index cb48294..7843860 100644 --- a/core/leras/layers/MsSsim.py +++ b/core/leras/layers/MsSsim.py @@ -24,6 +24,7 @@ class MsSsim(nn.LayerBase): ms_ssim_val = tf.image.ssim_multiscale(y_true_t, y_pred_t, max_val, power_factors=self.power_factors, filter_size=self.kernel_size) # ssim_multiscale returns values in range [0, 1] (where 1 is completely identical) # subtract from 1 to get loss - return 1.0 - ms_ssim_val + return (1.0 - ms_ssim_val) / 2.0 + nn.MsSsim = MsSsim From cce3061dca6e84d6979d2bc19973cbe1f232fc7e Mon Sep 17 00:00:00 2001 From: jh Date: Wed, 17 Mar 2021 11:41:46 -0700 Subject: [PATCH 32/44] fixme --- models/Model_SAEHD/Model.py | 20 ++++++++++++-------- 1 file changed, 12 insertions(+), 8 deletions(-) diff --git a/models/Model_SAEHD/Model.py b/models/Model_SAEHD/Model.py index 67dbe95..8de5bce 100644 --- a/models/Model_SAEHD/Model.py +++ b/models/Model_SAEHD/Model.py @@ -450,10 +450,13 @@ Examples: df, liae, df-d, df-ud, liae-ud, ... else: gpu_src_loss += tf.reduce_mean ( 300*tf.abs ( gpu_target_src*gpu_target_part_mask - gpu_pred_src_src*gpu_target_part_mask ), axis=[1,2,3]) - if self.options['ms_ssim_loss']: - gpu_src_loss += tf.reduce_mean ( 10*nn.MsSsim(resolution)(gpu_target_srcm, gpu_pred_src_srcm, max_val=1.0)) - else: - gpu_src_loss += tf.reduce_mean ( 10*tf.square( gpu_target_srcm - gpu_pred_src_srcm ),axis=[1,2,3] ) + # FIXME + # if self.options['ms_ssim_loss']: + # gpu_src_loss += tf.reduce_mean ( 10*nn.MsSsim(resolution)(gpu_target_srcm, gpu_pred_src_srcm, max_val=1.0)) + # else: + print('gpu_target_srcm.shape:', gpu_target_srcm.shape) + print('gpu_pred_src_srcm.shape:', gpu_pred_src_srcm.shape) + gpu_src_loss += tf.reduce_mean ( 10*tf.square( gpu_target_srcm - gpu_pred_src_srcm ),axis=[1,2,3] ) face_style_power = self.options['face_style_power'] / 100.0 if face_style_power != 0 and not self.pretrain: @@ -488,10 +491,11 @@ Examples: df, liae, df-d, df-ud, liae-ud, ... else: gpu_dst_loss += tf.reduce_mean ( 300*tf.abs ( gpu_target_dst*gpu_target_part_mask - gpu_pred_dst_dst*gpu_target_part_mask ), axis=[1,2,3]) - if self.options['ms_ssim_loss']: - gpu_dst_loss += tf.reduce_mean ( 10*nn.MsSsim(resolution)(gpu_target_dstm, gpu_pred_dst_dstm, max_val=1.0)) - else: - gpu_dst_loss += tf.reduce_mean ( 10*tf.square( gpu_target_dstm - gpu_pred_dst_dstm ),axis=[1,2,3] ) + # FIXME + # if self.options['ms_ssim_loss']: + # gpu_dst_loss += tf.reduce_mean ( 10*nn.MsSsim(resolution)(gpu_target_dstm, gpu_pred_dst_dstm, max_val=1.0)) + # else: + gpu_dst_loss += tf.reduce_mean ( 10*tf.square( gpu_target_dstm - gpu_pred_dst_dstm ),axis=[1,2,3] ) gpu_src_losses += [gpu_src_loss] gpu_dst_losses += [gpu_dst_loss] From bcffcde97f5e75b36d6ad18e021240b26c7e9c5d Mon Sep 17 00:00:00 2001 From: jh Date: Wed, 17 Mar 2021 11:53:38 -0700 Subject: [PATCH 33/44] debug --- models/Model_SAEHD/Model.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/models/Model_SAEHD/Model.py b/models/Model_SAEHD/Model.py index 8de5bce..a01836d 100644 --- a/models/Model_SAEHD/Model.py +++ b/models/Model_SAEHD/Model.py @@ -446,6 +446,16 @@ Examples: df, liae, df-d, df-ud, liae-ud, ... gpu_target_part_mask = gpu_target_srcm_mouth if self.options['ms_ssim_loss']: + # FIXME + ms_ssim_loss = nn.MsSsim(resolution, kernel_size=5)(gpu_target_src*gpu_target_part_mask, gpu_pred_src_src*gpu_target_part_mask, max_val=1.0) + print('ms_ssim_loss.shape: ', ms_ssim_loss.shape) + ms_ssim_mean_loss = tf.reduce_mean(300 * ms_ssim_loss) + print('ms_ssim_mean_loss.shape: ', ms_ssim_mean_loss.shape) + abs_loss = tf.abs ( gpu_target_src*gpu_target_part_mask - gpu_pred_src_src*gpu_target_part_mask ) + print('abs_loss.shape: ', abs_loss.shape) + abs_mean_loss = tf.reduce_mean(300 * abs_loss, axis=[1,2,3]) + print('abs_mean_loss.shape: ', abs_mean_loss.shape) + gpu_src_loss += tf.reduce_mean ( 300*nn.MsSsim(resolution, kernel_size=5)(gpu_target_src*gpu_target_part_mask, gpu_pred_src_src*gpu_target_part_mask, max_val=1.0)) else: gpu_src_loss += tf.reduce_mean ( 300*tf.abs ( gpu_target_src*gpu_target_part_mask - gpu_pred_src_src*gpu_target_part_mask ), axis=[1,2,3]) From e9ae398e8e33192733e96c3405e04e3c46ab466a Mon Sep 17 00:00:00 2001 From: jh Date: Wed, 17 Mar 2021 12:03:04 -0700 Subject: [PATCH 34/44] test --- models/Model_SAEHD/Model.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/models/Model_SAEHD/Model.py b/models/Model_SAEHD/Model.py index a01836d..be470df 100644 --- a/models/Model_SAEHD/Model.py +++ b/models/Model_SAEHD/Model.py @@ -455,8 +455,9 @@ Examples: df, liae, df-d, df-ud, liae-ud, ... print('abs_loss.shape: ', abs_loss.shape) abs_mean_loss = tf.reduce_mean(300 * abs_loss, axis=[1,2,3]) print('abs_mean_loss.shape: ', abs_mean_loss.shape) + gpu_src_loss += 300*nn.MsSsim(resolution, kernel_size=5)(gpu_target_src*gpu_target_part_mask, gpu_pred_src_src*gpu_target_part_mask, max_val=1.0) - gpu_src_loss += tf.reduce_mean ( 300*nn.MsSsim(resolution, kernel_size=5)(gpu_target_src*gpu_target_part_mask, gpu_pred_src_src*gpu_target_part_mask, max_val=1.0)) + # gpu_src_loss += tf.reduce_mean ( 300*nn.MsSsim(resolution, kernel_size=5)(gpu_target_src*gpu_target_part_mask, gpu_pred_src_src*gpu_target_part_mask, max_val=1.0)) else: gpu_src_loss += tf.reduce_mean ( 300*tf.abs ( gpu_target_src*gpu_target_part_mask - gpu_pred_src_src*gpu_target_part_mask ), axis=[1,2,3]) From 7eafa4edbd8dbf90d34bd1ec076f58ff2a6e129a Mon Sep 17 00:00:00 2001 From: jh Date: Wed, 17 Mar 2021 12:09:09 -0700 Subject: [PATCH 35/44] test --- models/Model_SAEHD/Model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/models/Model_SAEHD/Model.py b/models/Model_SAEHD/Model.py index be470df..2c706af 100644 --- a/models/Model_SAEHD/Model.py +++ b/models/Model_SAEHD/Model.py @@ -428,7 +428,7 @@ Examples: df, liae, df-d, df-ud, liae-ud, ... gpu_psd_target_dst_style_anti_masked = gpu_pred_src_dst*(1.0 - gpu_target_dstm_style_blur) if self.options['ms_ssim_loss']: - gpu_src_loss = tf.reduce_mean ( 10*nn.MsSsim(resolution)(gpu_target_src_masked_opt, gpu_pred_src_src_masked_opt, max_val=1.0)) + gpu_src_loss = 10 * nn.MsSsim(resolution)(gpu_target_src_masked_opt, gpu_pred_src_src_masked_opt, max_val=1.0) else: 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]) From 40924b5f74dc759fab39df1ccaef4a0a2306ba32 Mon Sep 17 00:00:00 2001 From: jh Date: Wed, 17 Mar 2021 12:13:28 -0700 Subject: [PATCH 36/44] test --- models/Model_SAEHD/Model.py | 24 +++++------------------- 1 file changed, 5 insertions(+), 19 deletions(-) diff --git a/models/Model_SAEHD/Model.py b/models/Model_SAEHD/Model.py index 2c706af..6ae3fb5 100644 --- a/models/Model_SAEHD/Model.py +++ b/models/Model_SAEHD/Model.py @@ -446,28 +446,14 @@ Examples: df, liae, df-d, df-ud, liae-ud, ... gpu_target_part_mask = gpu_target_srcm_mouth if self.options['ms_ssim_loss']: - # FIXME - ms_ssim_loss = nn.MsSsim(resolution, kernel_size=5)(gpu_target_src*gpu_target_part_mask, gpu_pred_src_src*gpu_target_part_mask, max_val=1.0) - print('ms_ssim_loss.shape: ', ms_ssim_loss.shape) - ms_ssim_mean_loss = tf.reduce_mean(300 * ms_ssim_loss) - print('ms_ssim_mean_loss.shape: ', ms_ssim_mean_loss.shape) - abs_loss = tf.abs ( gpu_target_src*gpu_target_part_mask - gpu_pred_src_src*gpu_target_part_mask ) - print('abs_loss.shape: ', abs_loss.shape) - abs_mean_loss = tf.reduce_mean(300 * abs_loss, axis=[1,2,3]) - print('abs_mean_loss.shape: ', abs_mean_loss.shape) - gpu_src_loss += 300*nn.MsSsim(resolution, kernel_size=5)(gpu_target_src*gpu_target_part_mask, gpu_pred_src_src*gpu_target_part_mask, max_val=1.0) - - # gpu_src_loss += tf.reduce_mean ( 300*nn.MsSsim(resolution, kernel_size=5)(gpu_target_src*gpu_target_part_mask, gpu_pred_src_src*gpu_target_part_mask, max_val=1.0)) + gpu_src_loss += 300 * nn.MsSsim(resolution, kernel_size=5)(gpu_target_src*gpu_target_part_mask, gpu_pred_src_src*gpu_target_part_mask, max_val=1.0) else: gpu_src_loss += tf.reduce_mean ( 300*tf.abs ( gpu_target_src*gpu_target_part_mask - gpu_pred_src_src*gpu_target_part_mask ), axis=[1,2,3]) - # FIXME - # if self.options['ms_ssim_loss']: - # gpu_src_loss += tf.reduce_mean ( 10*nn.MsSsim(resolution)(gpu_target_srcm, gpu_pred_src_srcm, max_val=1.0)) - # else: - print('gpu_target_srcm.shape:', gpu_target_srcm.shape) - print('gpu_pred_src_srcm.shape:', gpu_pred_src_srcm.shape) - gpu_src_loss += tf.reduce_mean ( 10*tf.square( gpu_target_srcm - gpu_pred_src_srcm ),axis=[1,2,3] ) + if self.options['ms_ssim_loss']: + gpu_src_loss += 10 * nn.MsSsim(resolution)(gpu_target_srcm, gpu_pred_src_srcm, max_val=1.0) + else: + gpu_src_loss += tf.reduce_mean ( 10*tf.square( gpu_target_srcm - gpu_pred_src_srcm ),axis=[1,2,3] ) face_style_power = self.options['face_style_power'] / 100.0 if face_style_power != 0 and not self.pretrain: From 4e088645316978fad82b10c917cce8ea6ee483a7 Mon Sep 17 00:00:00 2001 From: jh Date: Wed, 17 Mar 2021 12:16:03 -0700 Subject: [PATCH 37/44] test --- models/Model_SAEHD/Model.py | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/models/Model_SAEHD/Model.py b/models/Model_SAEHD/Model.py index 6ae3fb5..3f4864a 100644 --- a/models/Model_SAEHD/Model.py +++ b/models/Model_SAEHD/Model.py @@ -465,7 +465,7 @@ Examples: df, liae, df-d, df-ud, liae-ud, ... gpu_src_loss += tf.reduce_mean( (10*bg_style_power)*tf.square(gpu_psd_target_dst_style_anti_masked - gpu_target_dst_style_anti_masked), axis=[1,2,3] ) if self.options['ms_ssim_loss']: - gpu_dst_loss = tf.reduce_mean ( 10*nn.MsSsim(resolution)(gpu_target_dst_masked_opt, gpu_pred_dst_dst_masked_opt, max_val=1.0)) + gpu_dst_loss = 10*nn.MsSsim(resolution)(gpu_target_dst_masked_opt, gpu_pred_dst_dst_masked_opt, max_val=1.0) else: 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]) @@ -484,15 +484,14 @@ Examples: df, liae, df-d, df-ud, liae-ud, ... gpu_target_part_mask = gpu_target_dstm_mouth if self.options['ms_ssim_loss']: - gpu_dst_loss += tf.reduce_mean ( 300*nn.MsSsim(resolution, kernel_size=5)(gpu_target_dst*gpu_target_part_mask, gpu_pred_dst_dst*gpu_target_part_mask, max_val=1.0)) + gpu_dst_loss += 300 * nn.MsSsim(resolution, kernel_size=5)(gpu_target_dst*gpu_target_part_mask, gpu_pred_dst_dst*gpu_target_part_mask, max_val=1.0) else: gpu_dst_loss += tf.reduce_mean ( 300*tf.abs ( gpu_target_dst*gpu_target_part_mask - gpu_pred_dst_dst*gpu_target_part_mask ), axis=[1,2,3]) - # FIXME - # if self.options['ms_ssim_loss']: - # gpu_dst_loss += tf.reduce_mean ( 10*nn.MsSsim(resolution)(gpu_target_dstm, gpu_pred_dst_dstm, max_val=1.0)) - # else: - gpu_dst_loss += tf.reduce_mean ( 10*tf.square( gpu_target_dstm - gpu_pred_dst_dstm ),axis=[1,2,3] ) + if self.options['ms_ssim_loss']: + gpu_dst_loss += 10 * nn.MsSsim(resolution)(gpu_target_dstm, gpu_pred_dst_dstm, max_val=1.0) + else: + gpu_dst_loss += tf.reduce_mean ( 10*tf.square( gpu_target_dstm - gpu_pred_dst_dstm ),axis=[1,2,3] ) gpu_src_losses += [gpu_src_loss] gpu_dst_losses += [gpu_dst_loss] From cd6f820e9fc7bb36f39fb9b8b11a3015020c7625 Mon Sep 17 00:00:00 2001 From: jh Date: Wed, 17 Mar 2021 12:16:14 -0700 Subject: [PATCH 38/44] formatting --- models/Model_SAEHD/Model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/models/Model_SAEHD/Model.py b/models/Model_SAEHD/Model.py index 3f4864a..08bce7a 100644 --- a/models/Model_SAEHD/Model.py +++ b/models/Model_SAEHD/Model.py @@ -465,7 +465,7 @@ Examples: df, liae, df-d, df-ud, liae-ud, ... gpu_src_loss += tf.reduce_mean( (10*bg_style_power)*tf.square(gpu_psd_target_dst_style_anti_masked - gpu_target_dst_style_anti_masked), axis=[1,2,3] ) if self.options['ms_ssim_loss']: - gpu_dst_loss = 10*nn.MsSsim(resolution)(gpu_target_dst_masked_opt, gpu_pred_dst_dst_masked_opt, max_val=1.0) + gpu_dst_loss = 10 * nn.MsSsim(resolution)(gpu_target_dst_masked_opt, gpu_pred_dst_dst_masked_opt, max_val=1.0) else: 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]) From 5a8d9f66f252cc0d180cc8532c16b5ece386d147 Mon Sep 17 00:00:00 2001 From: jh Date: Wed, 17 Mar 2021 12:19:10 -0700 Subject: [PATCH 39/44] fixes loss range --- core/leras/layers/MsSsim.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/core/leras/layers/MsSsim.py b/core/leras/layers/MsSsim.py index 7843860..d45a599 100644 --- a/core/leras/layers/MsSsim.py +++ b/core/leras/layers/MsSsim.py @@ -24,7 +24,7 @@ class MsSsim(nn.LayerBase): ms_ssim_val = tf.image.ssim_multiscale(y_true_t, y_pred_t, max_val, power_factors=self.power_factors, filter_size=self.kernel_size) # ssim_multiscale returns values in range [0, 1] (where 1 is completely identical) # subtract from 1 to get loss - return (1.0 - ms_ssim_val) / 2.0 + return 1.0 - ms_ssim_val nn.MsSsim = MsSsim From 5dbbf4767cb6d03cb6cfec3d60c90ea8011548d6 Mon Sep 17 00:00:00 2001 From: jh Date: Wed, 17 Mar 2021 12:22:36 -0700 Subject: [PATCH 40/44] test hard placement --- core/leras/nn.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/core/leras/nn.py b/core/leras/nn.py index 8ac437b..d9c5850 100644 --- a/core/leras/nn.py +++ b/core/leras/nn.py @@ -112,7 +112,7 @@ class nn(): config = tf.ConfigProto(device_count={'GPU': 0}) else: nn.tf_default_device = "/GPU:0" - config = tf.ConfigProto(allow_soft_placement=True) + config = tf.ConfigProto() config.gpu_options.visible_device_list = ','.join([str(device.index) for device in device_config.devices]) config.gpu_options.force_gpu_compatible = True From ecb3436c1b1843b1b55448e1e1f813bb556e7bc8 Mon Sep 17 00:00:00 2001 From: jh Date: Wed, 17 Mar 2021 12:26:02 -0700 Subject: [PATCH 41/44] restore soft placement --- core/leras/nn.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/core/leras/nn.py b/core/leras/nn.py index d9c5850..8ac437b 100644 --- a/core/leras/nn.py +++ b/core/leras/nn.py @@ -112,7 +112,7 @@ class nn(): config = tf.ConfigProto(device_count={'GPU': 0}) else: nn.tf_default_device = "/GPU:0" - config = tf.ConfigProto() + config = tf.ConfigProto(allow_soft_placement=True) config.gpu_options.visible_device_list = ','.join([str(device.index) for device in device_config.devices]) config.gpu_options.force_gpu_compatible = True From c3a4e4159e6fe6a345f469ddce1dd8c262515e82 Mon Sep 17 00:00:00 2001 From: jh Date: Wed, 17 Mar 2021 12:37:06 -0700 Subject: [PATCH 42/44] rm commented out code --- core/leras/ops/__init__.py | 45 ++------------------------------------ 1 file changed, 2 insertions(+), 43 deletions(-) diff --git a/core/leras/ops/__init__.py b/core/leras/ops/__init__.py index f34461f..500a22a 100644 --- a/core/leras/ops/__init__.py +++ b/core/leras/ops/__init__.py @@ -307,47 +307,6 @@ def dssim(img1,img2, max_val, filter_size=11, filter_sigma=1.5, k1=0.01, k2=0.03 nn.dssim = dssim - -# def ms_ssim(img1, img2, resolution, kernel_size=11, k1=0.01, k2=0.03, max_value=1.0, -# power_factors=(0.0448, 0.2856, 0.3001, 0.2363, 0.1333)): -# -# # restrict mssim factors to those greater/equal to kernel size -# power_factors = [power_factors[i] for i in range(len(power_factors)) if resolution//(2**i) >= kernel_size] -# -# # normalize power factors if reduced because of size -# if sum(power_factors) < 1.0: -# power_factors = [x/sum(power_factors) for x in power_factors] -# -# img_dtype = img1.dtype -# if img_dtype != img2.dtype: -# raise ValueError("img1.dtype != img2.dtype") -# -# if img_dtype != tf.float32: -# img1 = tf.cast(img1, tf.float32) -# img2 = tf.cast(img2, tf.float32) -# -# # Transpose images from NCHW to NHWC -# img1_t = tf.transpose(img1, [0, 2, 3, 1]) -# img2_t = tf.transpose(img2, [0, 2, 3, 1]) -# -# def assign_device(op): -# if op.type != 'ListDiff': -# return '/gpu:0' -# else: -# return '/cpu:0' -# -# with tf.device(assign_device): -# ms_ssim_val = tf.image.ssim_multiscale(img1_t, img2_t, max_val=max_value, power_factors=power_factors, -# filter_size=kernel_size, k1=k1, k2=k2) -# ms_ssim_loss = (1.0 - ms_ssim_val) / 2.0 -# -# if img_dtype != tf.float32: -# ms_ssim_loss = tf.cast(ms_ssim_loss, img_dtype) -# return ms_ssim_loss -# -# nn.ms_ssim = ms_ssim - - def space_to_depth(x, size): if nn.data_format == "NHWC": # match NCHW version in order to switch data_format without problems @@ -426,7 +385,7 @@ def total_variation_mse(images): """ pixel_dif1 = images[:, 1:, :, :] - images[:, :-1, :, :] pixel_dif2 = images[:, :, 1:, :] - images[:, :, :-1, :] - + tot_var = ( tf.reduce_sum(tf.square(pixel_dif1), axis=[1,2,3]) + tf.reduce_sum(tf.square(pixel_dif2), axis=[1,2,3]) ) return tot_var @@ -441,4 +400,4 @@ def tf_suppress_lower_mean(t, eps=0.00001): q = tf.clip_by_value(q-t_mean_eps, 0, eps) q = q * (t/eps) return q -""" +""" \ No newline at end of file From 86493d386612ea612da6ff76b4b49f0b535396fd Mon Sep 17 00:00:00 2001 From: jh Date: Wed, 17 Mar 2021 14:18:27 -0700 Subject: [PATCH 43/44] docs: update docs for ms-ssim --- doc/features/ms-ssim/README.md | 42 +++++++++++++++++++++++++++++ doc/features/random-color/README.md | 5 +++- 2 files changed, 46 insertions(+), 1 deletion(-) create mode 100644 doc/features/ms-ssim/README.md diff --git a/doc/features/ms-ssim/README.md b/doc/features/ms-ssim/README.md new file mode 100644 index 0000000..c7a0476 --- /dev/null +++ b/doc/features/ms-ssim/README.md @@ -0,0 +1,42 @@ +# Multiscale SSIM (MS-SSIM) + +Allows you to train using the MS-SSIM (multiscale structural similarity index measure) as the main loss metric, +a perceptually more accurate measure of image quality than MSE (mean squared error). + +- [DESCRIPTION](#description) +- [USAGE](#usage) + +![](example.png) + +## DESCRIPTION + +[SSIM](https://en.wikipedia.org/wiki/Structural_similarity) is metric for comparing the perceptial quality of an image: +> SSIM is a perception-based model that considers image degradation as perceived change in structural information, +> while also incorporating important perceptual phenomena, including both luminance masking and contrast masking terms. +> [...] +> Structural information is the idea that the pixels have strong inter-dependencies especially when they are spatially +> close. These dependencies carry important information about the structure of the objects in the visual scene. +> Luminance masking is a phenomenon whereby image distortions (in this context) tend to be less visible in bright +> regions, while contrast masking is a phenomenon whereby distortions become less visible where there is significant +> activity or "texture" in the image. + +The current loss metric is a combination of SSIM (structural similarity index measure) and +[MSE](https://en.wikipedia.org/wiki/Mean_squared_error) (mean squared error). + +[Multiscale SSIM](https://en.wikipedia.org/wiki/Structural_similarity#Multi-Scale_SSIM) is a variant of SSIM that +improves upon SSIM by comparing the similarity at multiple scales (e.g.: full-size, half-size, 1/4 size, etc.) +By using MS-SSIM as our main loss metric, we should expect the image similarity to improve across each scale, improving +both the large scale and small scale detail of the predicted images. + +Original paper: [Wang, Zhou, Eero P. Simoncelli, and Alan C. Bovik. +"Multiscale structural similarity for image quality assessment." +Signals, Systems and Computers, 2004.](https://www.cns.nyu.edu/pub/eero/wang03b.pdf) + +## USAGE + +``` +[n] Use multiscale loss? ( y/n ?:help ) : y +``` + + + diff --git a/doc/features/random-color/README.md b/doc/features/random-color/README.md index 8d8298c..d1aeac1 100644 --- a/doc/features/random-color/README.md +++ b/doc/features/random-color/README.md @@ -19,4 +19,7 @@ maintaining the same `C*` (chroma, relative saturation). ## USAGE -`[n] Random color ( y/n ?:help ) : y` +``` +[n] Random color ( y/n ?:help ) : y +``` + From 0a1f35bf83e7bf28caf4cb820a45c29cad3e00ff Mon Sep 17 00:00:00 2001 From: jh Date: Wed, 24 Mar 2021 06:29:33 -0700 Subject: [PATCH 44/44] Updates ms-ssim --- models/Model_SAEHD/Model.py | 54 +++++++++++++++++-------------------- 1 file changed, 24 insertions(+), 30 deletions(-) diff --git a/models/Model_SAEHD/Model.py b/models/Model_SAEHD/Model.py index df9dee7..fcd4ef3 100644 --- a/models/Model_SAEHD/Model.py +++ b/models/Model_SAEHD/Model.py @@ -436,11 +436,11 @@ Examples: df, liae, df-d, df-ud, liae-ud, ... if self.options['ms_ssim_loss']: gpu_src_loss = 10 * nn.MsSsim(resolution)(gpu_target_src_masked_opt, gpu_pred_src_src_masked_opt, max_val=1.0) else: - 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]) + 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 or mouth_prio: @@ -451,23 +451,20 @@ Examples: df, liae, df-d, df-ud, liae-ud, ... elif mouth_prio: gpu_target_part_mask = gpu_target_srcm_mouth - if self.options['ms_ssim_loss']: - gpu_src_loss += 300 * nn.MsSsim(resolution, kernel_size=5)(gpu_target_src*gpu_target_part_mask, gpu_pred_src_src*gpu_target_part_mask, max_val=1.0) - else: gpu_src_loss += tf.reduce_mean ( 300*tf.abs ( gpu_target_src*gpu_target_part_mask - gpu_pred_src_src*gpu_target_part_mask ), axis=[1,2,3]) - if self.options['ms_ssim_loss']: - gpu_src_loss += 10 * nn.MsSsim(resolution)(gpu_target_srcm, gpu_pred_src_srcm, max_val=1.0) - else: gpu_src_loss += tf.reduce_mean ( 10*tf.square( gpu_target_srcm - gpu_pred_src_srcm ),axis=[1,2,3] ) if self.options['background_power'] > 0: bg_factor = self.options['background_power'] - if resolution < 256: - gpu_src_loss += bg_factor * tf.reduce_mean ( 10*nn.dssim(gpu_target_src_anti_masked, gpu_pred_src_src_anti_masked, max_val=1.0, filter_size=int(resolution/11.6)), axis=[1]) + if self.options['ms_ssim_loss']: + gpu_src_loss = 10 * nn.MsSsim(resolution)(gpu_target_src_anti_masked, gpu_pred_src_src_anti_masked, max_val=1.0) else: - gpu_src_loss += bg_factor * tf.reduce_mean ( 5*nn.dssim(gpu_target_src_anti_masked, gpu_pred_src_src_anti_masked, max_val=1.0, filter_size=int(resolution/11.6)), axis=[1]) - gpu_src_loss += bg_factor * tf.reduce_mean ( 5*nn.dssim(gpu_target_src_anti_masked, gpu_pred_src_src_anti_masked, max_val=1.0, filter_size=int(resolution/23.2)), axis=[1]) + if resolution < 256: + gpu_src_loss += bg_factor * tf.reduce_mean ( 10*nn.dssim(gpu_target_src_anti_masked, gpu_pred_src_src_anti_masked, max_val=1.0, filter_size=int(resolution/11.6)), axis=[1]) + else: + gpu_src_loss += bg_factor * tf.reduce_mean ( 5*nn.dssim(gpu_target_src_anti_masked, gpu_pred_src_src_anti_masked, max_val=1.0, filter_size=int(resolution/11.6)), axis=[1]) + gpu_src_loss += bg_factor * tf.reduce_mean ( 5*nn.dssim(gpu_target_src_anti_masked, gpu_pred_src_src_anti_masked, max_val=1.0, filter_size=int(resolution/23.2)), axis=[1]) gpu_src_loss += bg_factor * tf.reduce_mean ( 10*tf.square ( gpu_target_src_anti_masked - gpu_pred_src_src_anti_masked ), axis=[1,2,3]) face_style_power = self.options['face_style_power'] / 100.0 @@ -482,11 +479,11 @@ Examples: df, liae, df-d, df-ud, liae-ud, ... if self.options['ms_ssim_loss']: gpu_dst_loss = 10 * nn.MsSsim(resolution)(gpu_target_dst_masked_opt, gpu_pred_dst_dst_masked_opt, max_val=1.0) else: - 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]) + 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]) @@ -498,23 +495,20 @@ Examples: df, liae, df-d, df-ud, liae-ud, ... elif mouth_prio: gpu_target_part_mask = gpu_target_dstm_mouth - if self.options['ms_ssim_loss']: - gpu_dst_loss += 300 * nn.MsSsim(resolution, kernel_size=5)(gpu_target_dst*gpu_target_part_mask, gpu_pred_dst_dst*gpu_target_part_mask, max_val=1.0) - else: gpu_dst_loss += tf.reduce_mean ( 300*tf.abs ( gpu_target_dst*gpu_target_part_mask - gpu_pred_dst_dst*gpu_target_part_mask ), axis=[1,2,3]) if self.options['background_power'] > 0: bg_factor = self.options['background_power'] - if resolution < 256: - gpu_dst_loss += bg_factor * tf.reduce_mean ( 10*nn.dssim(gpu_target_dst_anti_masked, gpu_pred_dst_dst_anti_masked, max_val=1.0, filter_size=int(resolution/11.6)), axis=[1]) + if self.options['ms_ssim_loss']: + gpu_src_loss = 10 * nn.MsSsim(resolution)(gpu_target_dst_anti_masked, gpu_pred_dst_dst_anti_masked, max_val=1.0) else: - gpu_dst_loss += bg_factor * tf.reduce_mean ( 5*nn.dssim(gpu_target_dst_anti_masked, gpu_pred_dst_dst_anti_masked, max_val=1.0, filter_size=int(resolution/11.6)), axis=[1]) - gpu_dst_loss += bg_factor * tf.reduce_mean ( 5*nn.dssim(gpu_target_dst_anti_masked, gpu_pred_dst_dst_anti_masked, max_val=1.0, filter_size=int(resolution/23.2)), axis=[1]) + if resolution < 256: + gpu_dst_loss += bg_factor * tf.reduce_mean ( 10*nn.dssim(gpu_target_dst_anti_masked, gpu_pred_dst_dst_anti_masked, max_val=1.0, filter_size=int(resolution/11.6)), axis=[1]) + else: + gpu_dst_loss += bg_factor * tf.reduce_mean ( 5*nn.dssim(gpu_target_dst_anti_masked, gpu_pred_dst_dst_anti_masked, max_val=1.0, filter_size=int(resolution/11.6)), axis=[1]) + gpu_dst_loss += bg_factor * tf.reduce_mean ( 5*nn.dssim(gpu_target_dst_anti_masked, gpu_pred_dst_dst_anti_masked, max_val=1.0, filter_size=int(resolution/23.2)), axis=[1]) gpu_dst_loss += bg_factor * tf.reduce_mean ( 10*tf.square ( gpu_target_dst_anti_masked - gpu_pred_dst_dst_anti_masked ), axis=[1,2,3]) - if self.options['ms_ssim_loss']: - gpu_dst_loss += 10 * nn.MsSsim(resolution)(gpu_target_dstm, gpu_pred_dst_dstm, max_val=1.0) - else: gpu_dst_loss += tf.reduce_mean ( 10*tf.square( gpu_target_dstm - gpu_pred_dst_dstm ),axis=[1,2,3] ) gpu_src_losses += [gpu_src_loss]