mirror of
https://github.com/iperov/DeepFaceLab.git
synced 2025-07-16 10:03:41 -07:00
H64, H128, DF, LIAEF128: added pixel loss option.
This commit is contained in:
iperov
parent
af3dd59f67
commit
f8e63970d2
5 changed files with 52 additions and 34 deletions
|
|
@ -4,6 +4,7 @@ from nnlib import nnlib
|
|||
from models import ModelBase
|
||||
from facelib import FaceType
|
||||
from samples import *
|
||||
from utils.console_utils import *
|
||||
|
||||
class Model(ModelBase):
|
||||
|
||||
|
|
@ -11,6 +12,13 @@ class Model(ModelBase):
|
|||
decoder_srcH5 = 'decoder_src.h5'
|
||||
decoder_dstH5 = 'decoder_dst.h5'
|
||||
|
||||
#override
|
||||
def onInitializeOptions(self, is_first_run, ask_override):
|
||||
if is_first_run or ask_override:
|
||||
self.options['pixel_loss'] = self.options['pixel_loss'] = input_bool ("Use pixel loss? (y/n, ?:help skip: n/default ) : ", False, help_message="Default DSSIM loss good for initial understanding structure of faces. Use pixel loss after 30-40k epochs to enhance fine details and remove face jitter.")
|
||||
else:
|
||||
self.options['pixel_loss'] = self.options.get('pixel_loss', False)
|
||||
|
||||
#override
|
||||
def onInitialize(self, **in_options):
|
||||
exec(nnlib.import_all(), locals(), globals())
|
||||
|
|
@ -29,8 +37,8 @@ class Model(ModelBase):
|
|||
self.autoencoder_src = Model([ae_input_layer,mask_layer], self.decoder_src(self.encoder(ae_input_layer)))
|
||||
self.autoencoder_dst = Model([ae_input_layer,mask_layer], self.decoder_dst(self.encoder(ae_input_layer)))
|
||||
|
||||
self.autoencoder_src.compile(optimizer=Adam(lr=5e-5, beta_1=0.5, beta_2=0.999), loss=[DSSIMMaskLoss([mask_layer]), 'mse'] )
|
||||
self.autoencoder_dst.compile(optimizer=Adam(lr=5e-5, beta_1=0.5, beta_2=0.999), loss=[DSSIMMaskLoss([mask_layer]), 'mse'] )
|
||||
self.autoencoder_src.compile(optimizer=Adam(lr=5e-5, beta_1=0.5, beta_2=0.999), loss=[DSSIMMSEMaskLoss(mask_layer, is_mse=self.options['pixel_loss']), 'mse'] )
|
||||
self.autoencoder_dst.compile(optimizer=Adam(lr=5e-5, beta_1=0.5, beta_2=0.999), loss=[DSSIMMSEMaskLoss(mask_layer, is_mse=self.options['pixel_loss']), 'mse'] )
|
||||
|
||||
if self.is_training_mode:
|
||||
f = SampleProcessor.TypeFlags
|
||||
|
|
|
|||
|
|
@ -22,6 +22,11 @@ class Model(ModelBase):
|
|||
self.options.pop ('created_vram_gb')
|
||||
self.options['lighter_ae'] = self.options.get('lighter_ae', default_lighter_ae)
|
||||
|
||||
if is_first_run or ask_override:
|
||||
self.options['pixel_loss'] = self.options['pixel_loss'] = input_bool ("Use pixel loss? (y/n, ?:help skip: n/default ) : ", False, help_message="Default DSSIM loss good for initial understanding structure of faces. Use pixel loss after 30-40k epochs to enhance fine details and remove face jitter.")
|
||||
else:
|
||||
self.options['pixel_loss'] = self.options.get('pixel_loss', False)
|
||||
|
||||
#override
|
||||
def onInitialize(self, **in_options):
|
||||
exec(nnlib.import_all(), locals(), globals())
|
||||
|
|
@ -44,7 +49,7 @@ class Model(ModelBase):
|
|||
self.ae = Model([input_src_bgr,input_src_mask,input_dst_bgr,input_dst_mask], [rec_src_bgr, rec_src_mask, rec_dst_bgr, rec_dst_mask] )
|
||||
|
||||
self.ae.compile(optimizer=Adam(lr=5e-5, beta_1=0.5, beta_2=0.999),
|
||||
loss=[ DSSIMMaskLoss([input_src_mask]), 'mae', DSSIMMaskLoss([input_dst_mask]), 'mae' ] )
|
||||
loss=[ DSSIMMSEMaskLoss(input_src_mask, is_mse=self.options['pixel_loss']), 'mae', DSSIMMSEMaskLoss(input_dst_mask, is_mse=self.options['pixel_loss']), 'mae' ] )
|
||||
|
||||
self.src_view = K.function([input_src_bgr],[rec_src_bgr, rec_src_mask])
|
||||
self.dst_view = K.function([input_dst_bgr],[rec_dst_bgr, rec_dst_mask])
|
||||
|
|
|
|||
|
|
@ -22,6 +22,11 @@ class Model(ModelBase):
|
|||
self.options.pop ('created_vram_gb')
|
||||
self.options['lighter_ae'] = self.options.get('lighter_ae', default_lighter_ae)
|
||||
|
||||
if is_first_run or ask_override:
|
||||
self.options['pixel_loss'] = self.options['pixel_loss'] = input_bool ("Use pixel loss? (y/n, ?:help skip: n/default ) : ", False, help_message="Default DSSIM loss good for initial understanding structure of faces. Use pixel loss after 30-40k epochs to enhance fine details and remove face jitter.")
|
||||
else:
|
||||
self.options['pixel_loss'] = self.options.get('pixel_loss', False)
|
||||
|
||||
#override
|
||||
def onInitialize(self, **in_options):
|
||||
exec(nnlib.import_all(), locals(), globals())
|
||||
|
|
@ -45,8 +50,7 @@ class Model(ModelBase):
|
|||
|
||||
self.ae = Model([input_src_bgr,input_src_mask,input_dst_bgr,input_dst_mask], [rec_src_bgr, rec_src_mask, rec_dst_bgr, rec_dst_mask] )
|
||||
|
||||
self.ae.compile(optimizer=Adam(lr=5e-5, beta_1=0.5, beta_2=0.999),
|
||||
loss=[ DSSIMMaskLoss([input_src_mask]), 'mae', DSSIMMaskLoss([input_dst_mask]), 'mae' ] )
|
||||
self.ae.compile(optimizer=Adam(lr=5e-5, beta_1=0.5, beta_2=0.999), loss=[ DSSIMMSEMaskLoss(input_src_mask, is_mse=self.options['pixel_loss']), 'mae', DSSIMMSEMaskLoss(input_dst_mask, is_mse=self.options['pixel_loss']), 'mae' ] )
|
||||
|
||||
self.src_view = K.function([input_src_bgr],[rec_src_bgr, rec_src_mask])
|
||||
self.dst_view = K.function([input_dst_bgr],[rec_dst_bgr, rec_dst_mask])
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ from nnlib import nnlib
|
|||
from models import ModelBase
|
||||
from facelib import FaceType
|
||||
from samples import *
|
||||
from utils.console_utils import *
|
||||
|
||||
class Model(ModelBase):
|
||||
|
||||
|
|
@ -12,6 +13,13 @@ class Model(ModelBase):
|
|||
inter_BH5 = 'inter_B.h5'
|
||||
inter_ABH5 = 'inter_AB.h5'
|
||||
|
||||
#override
|
||||
def onInitializeOptions(self, is_first_run, ask_override):
|
||||
if is_first_run or ask_override:
|
||||
self.options['pixel_loss'] = self.options['pixel_loss'] = input_bool ("Use pixel loss? (y/n, ?:help skip: n/default ) : ", False, help_message="Default DSSIM loss good for initial understanding structure of faces. Use pixel loss after 30-40k epochs to enhance fine details and remove face jitter.")
|
||||
else:
|
||||
self.options['pixel_loss'] = self.options.get('pixel_loss', False)
|
||||
|
||||
#override
|
||||
def onInitialize(self, **in_options):
|
||||
exec(nnlib.import_all(), locals(), globals())
|
||||
|
|
@ -34,8 +42,8 @@ class Model(ModelBase):
|
|||
self.autoencoder_src = Model([ae_input_layer,mask_layer], self.decoder(Concatenate()([AB, AB])) )
|
||||
self.autoencoder_dst = Model([ae_input_layer,mask_layer], self.decoder(Concatenate()([B, AB])) )
|
||||
|
||||
self.autoencoder_src.compile(optimizer=Adam(lr=5e-5, beta_1=0.5, beta_2=0.999), loss=[DSSIMMaskLoss([mask_layer]), 'mse'] )
|
||||
self.autoencoder_dst.compile(optimizer=Adam(lr=5e-5, beta_1=0.5, beta_2=0.999), loss=[DSSIMMaskLoss([mask_layer]), 'mse'] )
|
||||
self.autoencoder_src.compile(optimizer=Adam(lr=5e-5, beta_1=0.5, beta_2=0.999), loss=[DSSIMMSEMaskLoss(mask_layer, is_mse=self.options['pixel_loss']), 'mse'] )
|
||||
self.autoencoder_dst.compile(optimizer=Adam(lr=5e-5, beta_1=0.5, beta_2=0.999), loss=[DSSIMMSEMaskLoss(mask_layer, is_mse=self.options['pixel_loss']), 'mse'] )
|
||||
|
||||
if self.is_training_mode:
|
||||
f = SampleProcessor.TypeFlags
|
||||
|
|
|
|||
|
|
@ -37,7 +37,7 @@ class nnlib(object):
|
|||
modelify = None
|
||||
ReflectionPadding2D = None
|
||||
DSSIMLoss = None
|
||||
DSSIMMaskLoss = None
|
||||
DSSIMMSEMaskLoss = None
|
||||
PixelShuffler = None
|
||||
SubpixelUpscaler = None
|
||||
AddUniformNoise = None
|
||||
|
|
@ -101,7 +101,7 @@ Adam = keras.optimizers.Adam
|
|||
modelify = nnlib.modelify
|
||||
ReflectionPadding2D = nnlib.ReflectionPadding2D
|
||||
DSSIMLoss = nnlib.DSSIMLoss
|
||||
DSSIMMaskLoss = nnlib.DSSIMMaskLoss
|
||||
DSSIMMSEMaskLoss = nnlib.DSSIMMSEMaskLoss
|
||||
PixelShuffler = nnlib.PixelShuffler
|
||||
SubpixelUpscaler = nnlib.SubpixelUpscaler
|
||||
AddUniformNoise = nnlib.AddUniformNoise
|
||||
|
|
@ -417,6 +417,7 @@ NLayerDiscriminator = nnlib.NLayerDiscriminator
|
|||
tf = nnlib.tf
|
||||
keras = nnlib.keras
|
||||
K = keras.backend
|
||||
exec (nnlib.code_import_tf, locals(), globals())
|
||||
|
||||
def modelify(model_functor):
|
||||
def func(tensor):
|
||||
|
|
@ -451,29 +452,21 @@ NLayerDiscriminator = nnlib.NLayerDiscriminator
|
|||
return (1.0 - tf.image.ssim ((y_true/2+0.5), (y_pred/2+0.5), 1.0)) / 2.0
|
||||
nnlib.DSSIMLoss = DSSIMLoss
|
||||
|
||||
class DSSIMMaskLoss(object):
|
||||
def __init__(self, mask_list, is_tanh=False):
|
||||
self.mask_list = mask_list
|
||||
self.is_tanh = is_tanh
|
||||
class DSSIMMSEMaskLoss(object):
|
||||
def __init__(self, mask, is_mse=False):
|
||||
self.mask = mask
|
||||
self.is_mse = is_mse
|
||||
|
||||
def __call__(self,y_true, y_pred):
|
||||
total_loss = None
|
||||
for mask in self.mask_list:
|
||||
|
||||
if not self.is_tanh:
|
||||
loss = (1.0 - (tf.image.ssim (y_true*mask, y_pred*mask, 1.0))) / 2.0
|
||||
mask = self.mask
|
||||
if self.is_mse:
|
||||
blur_mask = tf_gaussian_blur(max(1, mask.get_shape().as_list()[1] // 32))(mask)
|
||||
return K.mean ( 100*K.square( y_true*blur_mask - y_pred*blur_mask ) )
|
||||
else:
|
||||
loss = (1.0 - tf.image.ssim ( (y_true/2+0.5)*(mask/2+0.5), (y_pred/2+0.5)*(mask/2+0.5), 1.0)) / 2.0
|
||||
|
||||
loss = K.cast (loss, K.floatx())
|
||||
|
||||
if total_loss is None:
|
||||
total_loss = loss
|
||||
else:
|
||||
total_loss += loss
|
||||
|
||||
return total_loss
|
||||
nnlib.DSSIMMaskLoss = DSSIMMaskLoss
|
||||
return (1.0 - (tf.image.ssim (y_true*mask, y_pred*mask, 1.0))) / 2.0
|
||||
nnlib.DSSIMMSEMaskLoss = DSSIMMSEMaskLoss
|
||||
|
||||
class PixelShuffler(keras.layers.Layer):
|
||||
def __init__(self, size=(2, 2), data_format=None, **kwargs):
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue