From 6c5c248f22857c28e548b6241984ecfeec1634a5 Mon Sep 17 00:00:00 2001 From: Jan Date: Thu, 9 Dec 2021 19:45:40 +0100 Subject: [PATCH] removed v3 gan --- core/leras/models/PatchDiscriminator.py | 129 +----------------------- models/Model_AMP/Model.py | 30 ++---- models/Model_SAEHD/Model.py | 35 +++---- 3 files changed, 23 insertions(+), 171 deletions(-) diff --git a/core/leras/models/PatchDiscriminator.py b/core/leras/models/PatchDiscriminator.py index 66ab02e..6cddb7f 100644 --- a/core/leras/models/PatchDiscriminator.py +++ b/core/leras/models/PatchDiscriminator.py @@ -76,134 +76,9 @@ class PatchDiscriminator(nn.ModelBase): nn.PatchDiscriminator = PatchDiscriminator - class UNetPatchDiscriminator(nn.ModelBase): """ Inspired by https://arxiv.org/abs/2002.12655 "A U-Net Based Discriminator for Generative Adversarial Networks" - Based on iperov commit 11add4cd4f5a61df26a8659f4cc5c8d9467bf5f8 from Jan 3, 2021 with added fp16 option - """ - def calc_receptive_field_size(self, layers): - """ - result the same as https://fomoro.com/research/article/receptive-field-calculatorindex.html - """ - rf = 0 - ts = 1 - for i, (k, s) in enumerate(layers): - if i == 0: - rf = k - else: - rf += (k-1)*ts - ts *= s - return rf - - def find_archi(self, target_patch_size, max_layers=9): - """ - Find the best configuration of layers using only 3x3 convs for target patch size - """ - s = {} - for layers_count in range(1,max_layers+1): - val = 1 << (layers_count-1) - while True: - val -= 1 - - layers = [] - sum_st = 0 - layers.append ( [3, 2]) - sum_st += 2 - for i in range(layers_count-1): - st = 1 + (1 if val & (1 << i) !=0 else 0 ) - layers.append ( [3, st ]) - sum_st += st - - rf = self.calc_receptive_field_size(layers) - - s_rf = s.get(rf, None) - if s_rf is None: - s[rf] = (layers_count, sum_st, layers) - else: - if layers_count < s_rf[0] or \ - ( layers_count == s_rf[0] and sum_st > s_rf[1] ): - s[rf] = (layers_count, sum_st, layers) - - if val == 0: - break - - x = sorted(list(s.keys())) - q=x[np.abs(np.array(x)-target_patch_size).argmin()] - return s[q][2] - - def on_build(self, patch_size, in_ch, base_ch = 16, use_fp16 = False): - self.use_fp16 = use_fp16 - conv_dtype = tf.float16 if use_fp16 else tf.float32 - - class ResidualBlock(nn.ModelBase): - def on_build(self, ch, kernel_size=3 ): - self.conv1 = nn.Conv2D( ch, ch, kernel_size=kernel_size, padding='SAME', dtype=conv_dtype) - self.conv2 = nn.Conv2D( ch, ch, kernel_size=kernel_size, padding='SAME', dtype=conv_dtype) - - def forward(self, inp): - x = self.conv1(inp) - x = tf.nn.leaky_relu(x, 0.2) - x = self.conv2(x) - x = tf.nn.leaky_relu(inp + x, 0.2) - return x - - prev_ch = in_ch - self.convs = [] - self.res1 = [] - self.res2 = [] - self.upconvs = [] - self.upres1 = [] - self.upres2 = [] - layers = self.find_archi(patch_size) - - level_chs = { i-1:v for i,v in enumerate([ min( base_ch * (2**i), 512 ) for i in range(len(layers)+1)]) } - - self.in_conv = nn.Conv2D( in_ch, level_chs[-1], kernel_size=1, padding='VALID', dtype=conv_dtype) - - for i, (kernel_size, strides) in enumerate(layers): - self.convs.append ( nn.Conv2D( level_chs[i-1], level_chs[i], kernel_size=kernel_size, strides=strides, padding='SAME', dtype=conv_dtype) ) - - self.res1.append ( ResidualBlock(level_chs[i]) ) - self.res2.append ( ResidualBlock(level_chs[i]) ) - - self.upconvs.insert (0, nn.Conv2DTranspose( level_chs[i]*(2 if i != len(layers)-1 else 1), level_chs[i-1], kernel_size=kernel_size, strides=strides, padding='SAME', dtype=conv_dtype) ) - - self.upres1.insert (0, ResidualBlock(level_chs[i-1]*2) ) - self.upres2.insert (0, ResidualBlock(level_chs[i-1]*2) ) - - self.out_conv = nn.Conv2D( level_chs[-1]*2, 1, kernel_size=1, padding='VALID', dtype=conv_dtype) - - self.center_out = nn.Conv2D( level_chs[len(layers)-1], 1, kernel_size=1, padding='VALID', dtype=conv_dtype) - self.center_conv = nn.Conv2D( level_chs[len(layers)-1], level_chs[len(layers)-1], kernel_size=1, padding='VALID', dtype=conv_dtype) - - - def forward(self, x): - x = tf.nn.leaky_relu( self.in_conv(x), 0.2 ) - - encs = [] - for conv, res1,res2 in zip(self.convs, self.res1, self.res2): - encs.insert(0, x) - x = tf.nn.leaky_relu( conv(x), 0.2 ) - x = res1(x) - x = res2(x) - - center_out, x = self.center_out(x), tf.nn.leaky_relu( self.center_conv(x), 0.2 ) - - for i, (upconv, enc, upres1, upres2 ) in enumerate(zip(self.upconvs, encs, self.upres1, self.upres2)): - x = tf.nn.leaky_relu( upconv(x), 0.2 ) - x = tf.concat( [enc, x], axis=nn.conv2d_ch_axis) - x = upres1(x) - x = upres2(x) - - return center_out, self.out_conv(x) - -nn.UNetPatchDiscriminator = UNetPatchDiscriminator - -class UNetPatchDiscriminatorV2(nn.ModelBase): - """ - Inspired by https://arxiv.org/abs/2002.12655 "A U-Net Based Discriminator for Generative Adversarial Networks" - Based on iperov commit 35877dbfd724c22040f421e93c1adbb7142e5b5d from Jul 14, 2021 """ def calc_receptive_field_size(self, layers): """ @@ -315,5 +190,5 @@ class UNetPatchDiscriminatorV2(nn.ModelBase): x = tf.cast(x, tf.float32) return center_out, x - -nn.UNetPatchDiscriminatorV2 = UNetPatchDiscriminatorV2 + +nn.UNetPatchDiscriminator = UNetPatchDiscriminator \ No newline at end of file diff --git a/models/Model_AMP/Model.py b/models/Model_AMP/Model.py index 501db23..d57ba7d 100644 --- a/models/Model_AMP/Model.py +++ b/models/Model_AMP/Model.py @@ -132,7 +132,6 @@ class AMPModel(ModelBase): 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.") default_gan_power = self.options['gan_power'] = self.load_or_def_option('gan_power', 0.0) - default_gan_version = self.options['gan_version'] = self.load_or_def_option('gan_version', 2) 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) default_gan_smoothing = self.options['gan_smoothing'] = self.load_or_def_option('gan_smoothing', 0.1) @@ -154,14 +153,12 @@ class AMPModel(ModelBase): if self.options['gan_power'] != 0.0: - self.options['gan_version'] = np.clip (io.input_int("GAN version", default_gan_version, add_info="2 or 3", help_message="Choose GAN version (v2: 7/16/2020, v3: 1/3/2021):"), 2, 3) - if self.options['gan_version'] == 3: - 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_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 + 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 self.options['gan_smoothing'] = np.clip ( io.input_number("GAN label smoothing", default_gan_smoothing, add_info="0 - 0.5", help_message="Uses soft labels with values slightly off from 0/1 for GAN, has a regularizing effect"), 0, 0.5) self.options['gan_noise'] = np.clip ( io.input_number("GAN noisy labels", default_gan_noise, add_info="0 - 0.5", help_message="Marks some images with the wrong label, helps prevent collapse"), 0, 0.5) @@ -381,12 +378,8 @@ class AMPModel(ModelBase): if self.is_training: if gan_power != 0: - if self.options['gan_version'] == 2: - self.GAN = nn.UNetPatchDiscriminatorV2(patch_size=resolution//16, in_ch=input_ch, name="D_src", use_fp16=self.options['use_fp16']) - self.model_filename_list += [ [self.GAN, 'D_src_v2.npy'] ] - else: - self.GAN = nn.UNetPatchDiscriminator(patch_size=self.options['gan_patch_size'], in_ch=input_ch, base_ch=self.options['gan_dims'], use_fp16=self.options['use_fp16'], name="D_src") - self.model_filename_list += [ [self.GAN, 'GAN.npy'] ] + self.GAN = nn.UNetPatchDiscriminator(patch_size=self.options['gan_patch_size'], in_ch=input_ch, base_ch=self.options['gan_dims'], use_fp16=self.options['use_fp16'], name="D_src") + self.model_filename_list += [ [self.GAN, 'GAN.npy'] ] # Initialize optimizers lr_modifier = self.options['lr_modifier'] @@ -411,14 +404,9 @@ class AMPModel(ModelBase): self.model_filename_list += [ (self.src_dst_opt, 'src_dst_opt.npy') ] if gan_power != 0: - if self.options['gan_version'] == 2: - self.GAN_opt = OptimizerClass(lr=lr, lr_dropout=lr_dropout, lr_cos=lr_cos, clipnorm=clipnorm, name='D_src_dst_opt') - self.GAN_opt.initialize_variables ( self.GAN.get_weights(), vars_on_cpu=optimizer_vars_on_cpu, lr_dropout_on_cpu=self.options['lr_dropout']=='cpu')#+self.D_src_x2.get_weights() - self.model_filename_list += [ (self.GAN_opt, 'D_src_v2_opt.npy') ] - else: - self.GAN_opt = OptimizerClass(lr=lr, lr_dropout=lr_dropout, lr_cos=lr_cos, clipnorm=clipnorm, name='GAN_opt') - self.GAN_opt.initialize_variables ( self.GAN.get_weights(), vars_on_cpu=optimizer_vars_on_cpu, lr_dropout_on_cpu=self.options['lr_dropout']=='cpu')#+self.D_src_x2.get_weights() - self.model_filename_list += [ (self.GAN_opt, 'GAN_opt.npy') ] + self.GAN_opt = OptimizerClass(lr=lr, lr_dropout=lr_dropout, lr_cos=lr_cos, clipnorm=clipnorm, name='GAN_opt') + self.GAN_opt.initialize_variables ( self.GAN.get_weights(), vars_on_cpu=optimizer_vars_on_cpu, lr_dropout_on_cpu=self.options['lr_dropout']=='cpu')#+self.D_src_x2.get_weights() + self.model_filename_list += [ (self.GAN_opt, 'GAN_opt.npy') ] if self.is_training: # Adjust batch size for multiple GPU diff --git a/models/Model_SAEHD/Model.py b/models/Model_SAEHD/Model.py index d79e92a..f3ff7e2 100644 --- a/models/Model_SAEHD/Model.py +++ b/models/Model_SAEHD/Model.py @@ -165,7 +165,6 @@ class SAEHDModel(ModelBase): self.options['blur_out_mask'] = io.input_bool ("Blur out mask", default_blur_out_mask, help_message='Blurs nearby area outside of applied face mask of training samples. The result is the background near the face is smoothed and less noticeable on swapped face. The exact xseg mask in src and dst faceset is required.') default_gan_power = self.options['gan_power'] = self.load_or_def_option('gan_power', 0.0) - default_gan_version = self.options['gan_version'] = self.load_or_def_option('gan_version', 2) 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) default_gan_smoothing = self.options['gan_smoothing'] = self.load_or_def_option('gan_smoothing', 0.1) @@ -194,17 +193,15 @@ class SAEHDModel(ModelBase): self.options['gan_power'] = np.clip ( io.input_number ("GAN power", default_gan_power, add_info="0.0 .. 10.0", help_message="Train the network in Generative Adversarial manner. Forces the neural network to learn small details of the face. Enable it only when the face is trained enough and don't disable. Typical value is 0.1"), 0.0, 10.0 ) if self.options['gan_power'] != 0.0: - self.options['gan_version'] = np.clip (io.input_int("GAN version", default_gan_version, add_info="2 or 3", help_message="Choose GAN version (v2: 7/16/2020, v3: 1/3/2021):"), 2, 3) - if self.options['gan_version'] == 3: - 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_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 + 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 - self.options['gan_smoothing'] = np.clip ( io.input_number("GAN label smoothing", default_gan_smoothing, add_info="0 - 0.5", help_message="Uses soft labels with values slightly off from 0/1 for GAN, has a regularizing effect"), 0, 0.5) - self.options['gan_noise'] = np.clip ( io.input_number("GAN noisy labels", default_gan_noise, add_info="0 - 0.5", help_message="Marks some images with the wrong label, helps prevent collapse"), 0, 0.5) + self.options['gan_smoothing'] = np.clip ( io.input_number("GAN label smoothing", default_gan_smoothing, add_info="0 - 0.5", help_message="Uses soft labels with values slightly off from 0/1 for GAN, has a regularizing effect"), 0, 0.5) + self.options['gan_noise'] = np.clip ( io.input_number("GAN noisy labels", default_gan_noise, add_info="0 - 0.5", help_message="Marks some images with the wrong label, helps prevent collapse"), 0, 0.5) 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 ) @@ -355,12 +352,8 @@ class SAEHDModel(ModelBase): if self.is_training: if gan_power != 0: - if self.options['gan_version'] == 2: - self.D_src = nn.UNetPatchDiscriminatorV2(patch_size=resolution//16, in_ch=input_ch, name="D_src", use_fp16=self.options['use_fp16']) - self.model_filename_list += [ [self.D_src, 'D_src_v2.npy'] ] - else: - self.D_src = nn.UNetPatchDiscriminator(patch_size=self.options['gan_patch_size'], in_ch=input_ch, base_ch=self.options['gan_dims'], use_fp16=self.options['use_fp16'], name="D_src") - self.model_filename_list += [ [self.D_src, 'GAN.npy'] ] + self.D_src = nn.UNetPatchDiscriminator(patch_size=self.options['gan_patch_size'], in_ch=input_ch, base_ch=self.options['gan_dims'], use_fp16=self.options['use_fp16'], name="D_src") + self.model_filename_list += [ [self.D_src, 'GAN.npy'] ] # Initialize optimizers lr_modifier = self.options['lr_modifier'] @@ -400,14 +393,10 @@ class SAEHDModel(ModelBase): self.model_filename_list += [ (self.D_code_opt, 'D_code_opt.npy') ] if gan_power != 0: - if self.options['gan_version'] == 2: - self.D_src_dst_opt = OptimizerClass(lr=lr, lr_dropout=lr_dropout, lr_cos=lr_cos, clipnorm=clipnorm, name='D_src_dst_opt') - self.D_src_dst_opt.initialize_variables ( self.D_src.get_weights(), vars_on_cpu=optimizer_vars_on_cpu, lr_dropout_on_cpu=self.options['lr_dropout']=='cpu')#+self.D_src_x2.get_weights() - self.model_filename_list += [ (self.D_src_dst_opt, 'D_src_v2_opt.npy') ] - else: - self.D_src_dst_opt = OptimizerClass(lr=lr, lr_dropout=lr_dropout, lr_cos=lr_cos, clipnorm=clipnorm, name='GAN_opt') - self.D_src_dst_opt.initialize_variables ( self.D_src.get_weights(), vars_on_cpu=optimizer_vars_on_cpu, lr_dropout_on_cpu=self.options['lr_dropout']=='cpu')#+self.D_src_x2.get_weights() - self.model_filename_list += [ (self.D_src_dst_opt, 'GAN_opt.npy') ] + self.D_src_dst_opt = OptimizerClass(lr=lr, lr_dropout=lr_dropout, lr_cos=lr_cos, clipnorm=clipnorm, name='GAN_opt') + self.D_src_dst_opt.initialize_variables ( self.D_src.get_weights(), vars_on_cpu=optimizer_vars_on_cpu, lr_dropout_on_cpu=self.options['lr_dropout']=='cpu')#+self.D_src_x2.get_weights() + self.model_filename_list += [ (self.D_src_dst_opt, 'GAN_opt.npy') ] + if self.is_training: # Adjust batch size for multiple GPU