mirror of
https://github.com/iperov/DeepFaceLab.git
synced 2025-07-06 21:12:07 -07:00
increased speed of train on H64, H128 models.
This commit is contained in:
iperov
parent
77640259fc
commit
612ef5155e
2 changed files with 159 additions and 134 deletions
|
|
@ -15,30 +15,35 @@ class Model(ModelBase):
|
|||
|
||||
#override
|
||||
def onInitialize(self, **in_options):
|
||||
tf = self.tf
|
||||
keras = self.keras
|
||||
K = keras.backend
|
||||
self.set_vram_batch_requirements( {1.5:2,2:2,3:4,4:8,5:16,6:32,7:32,8:32,9:48} )
|
||||
|
||||
ae_input_layer = self.keras.layers.Input(shape=(64, 64, 3))
|
||||
mask_layer = self.keras.layers.Input(shape=(64, 64, 1)) #same as output
|
||||
|
||||
self.encoder = self.Encoder(ae_input_layer, self.created_vram_gb)
|
||||
self.decoder_src = self.Decoder(self.created_vram_gb)
|
||||
self.decoder_dst = self.Decoder(self.created_vram_gb)
|
||||
|
||||
bgr_shape, mask_shape, self.encoder, self.decoder_src, self.decoder_dst = self.Build(self.created_vram_gb)
|
||||
if not self.is_first_run():
|
||||
self.encoder.load_weights (self.get_strpath_storage_for_file(self.encoderH5))
|
||||
self.decoder_src.load_weights (self.get_strpath_storage_for_file(self.decoder_srcH5))
|
||||
self.decoder_dst.load_weights (self.get_strpath_storage_for_file(self.decoder_dstH5))
|
||||
|
||||
input_src_bgr = self.keras.layers.Input(bgr_shape)
|
||||
input_src_mask = self.keras.layers.Input(mask_shape)
|
||||
input_dst_bgr = self.keras.layers.Input(bgr_shape)
|
||||
input_dst_mask = self.keras.layers.Input(mask_shape)
|
||||
|
||||
self.autoencoder_src = self.keras.models.Model([ae_input_layer,mask_layer], self.decoder_src(self.encoder(ae_input_layer)))
|
||||
self.autoencoder_dst = self.keras.models.Model([ae_input_layer,mask_layer], self.decoder_dst(self.encoder(ae_input_layer)))
|
||||
rec_src_bgr, rec_src_mask = self.decoder_src( self.encoder(input_src_bgr) )
|
||||
rec_dst_bgr, rec_dst_mask = self.decoder_dst( self.encoder(input_dst_bgr) )
|
||||
|
||||
self.ae = self.keras.models.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] )
|
||||
|
||||
if self.is_training_mode:
|
||||
self.autoencoder_src, self.autoencoder_dst = self.to_multi_gpu_model_if_possible ( [self.autoencoder_src, self.autoencoder_dst] )
|
||||
|
||||
optimizer = self.keras.optimizers.Adam(lr=5e-5, beta_1=0.5, beta_2=0.999)
|
||||
dssimloss = DSSIMMaskLossClass(self.tf)([mask_layer])
|
||||
self.autoencoder_src.compile(optimizer=optimizer, loss=[dssimloss, 'mae'])
|
||||
self.autoencoder_dst.compile(optimizer=optimizer, loss=[dssimloss, 'mae'])
|
||||
self.ae, = self.to_multi_gpu_model_if_possible ( [self.ae,] )
|
||||
|
||||
self.ae.compile(optimizer=self.keras.optimizers.Adam(lr=5e-5, beta_1=0.5, beta_2=0.999),
|
||||
loss=[ DSSIMMaskLossClass(self.tf)([input_src_mask]), 'mae', DSSIMMaskLossClass(self.tf)([input_dst_mask]), '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])
|
||||
|
||||
if self.is_training_mode:
|
||||
from models import TrainingDataGenerator
|
||||
|
|
@ -58,11 +63,11 @@ class Model(ModelBase):
|
|||
def onTrainOneEpoch(self, sample):
|
||||
warped_src, target_src, target_src_full_mask = sample[0]
|
||||
warped_dst, target_dst, target_dst_full_mask = sample[1]
|
||||
|
||||
loss_src = self.autoencoder_src.train_on_batch( [warped_src, target_src_full_mask], [target_src, target_src_full_mask] )
|
||||
loss_dst = self.autoencoder_dst.train_on_batch( [warped_dst, target_dst_full_mask], [target_dst, target_dst_full_mask] )
|
||||
|
||||
return ( ('loss_src', loss_src[0]), ('loss_dst', loss_dst[0]) )
|
||||
|
||||
total, loss_src_bgr, loss_src_mask, loss_dst_bgr, loss_dst_mask = self.ae.train_on_batch( [warped_src, target_src_full_mask, warped_dst, target_dst_full_mask], [target_src, target_src_full_mask, target_dst, target_dst_full_mask] )
|
||||
|
||||
return ( ('loss_src', loss_src_bgr), ('loss_dst', loss_dst_bgr) )
|
||||
|
||||
#override
|
||||
def onGetPreview(self, sample):
|
||||
|
|
@ -71,9 +76,9 @@ class Model(ModelBase):
|
|||
test_B = sample[1][1][0:4]
|
||||
test_B_m = sample[1][2][0:4]
|
||||
|
||||
AA, mAA = self.autoencoder_src.predict([test_A, test_A_m])
|
||||
AB, mAB = self.autoencoder_src.predict([test_B, test_B_m])
|
||||
BB, mBB = self.autoencoder_dst.predict([test_B, test_B_m])
|
||||
AA, mAA = self.src_view([test_A])
|
||||
AB, mAB = self.src_view([test_B])
|
||||
BB, mBB = self.dst_view([test_B])
|
||||
|
||||
mAA = np.repeat ( mAA, (3,), -1)
|
||||
mAB = np.repeat ( mAB, (3,), -1)
|
||||
|
|
@ -99,7 +104,7 @@ class Model(ModelBase):
|
|||
face_64_bgr = face[...,0:3]
|
||||
face_64_mask = np.expand_dims(face[...,3],-1)
|
||||
|
||||
x, mx = self.autoencoder_src.predict ( [ np.expand_dims(face_64_bgr,0), np.expand_dims(face_64_mask,0) ] )
|
||||
x, mx = self.src_view ( [ np.expand_dims(face_64_bgr,0) ] )
|
||||
x, mx = x[0], mx[0]
|
||||
|
||||
return np.concatenate ( (x,mx), -1 )
|
||||
|
|
@ -121,47 +126,54 @@ class Model(ModelBase):
|
|||
|
||||
return ConverterMasked(self.predictor_func, predictor_input_size=64, output_size=64, face_type=FaceType.HALF, **in_options)
|
||||
|
||||
def Encoder(self, input_layer, created_vram_gb):
|
||||
x = input_layer
|
||||
if created_vram_gb >= 4:
|
||||
x = conv(self.keras, x, 128)
|
||||
x = conv(self.keras, x, 256)
|
||||
x = conv(self.keras, x, 512)
|
||||
x = conv(self.keras, x, 1024)
|
||||
x = self.keras.layers.Dense(1024)(self.keras.layers.Flatten()(x))
|
||||
x = self.keras.layers.Dense(4 * 4 * 1024)(x)
|
||||
x = self.keras.layers.Reshape((4, 4, 1024))(x)
|
||||
x = upscale(self.keras, x, 512)
|
||||
else:
|
||||
x = conv(self.keras, x, 128 )
|
||||
x = conv(self.keras, x, 256 )
|
||||
x = conv(self.keras, x, 512 )
|
||||
x = conv(self.keras, x, 768 )
|
||||
x = self.keras.layers.Dense(512)(self.keras.layers.Flatten()(x))
|
||||
x = self.keras.layers.Dense(4 * 4 * 512)(x)
|
||||
x = self.keras.layers.Reshape((4, 4, 512))(x)
|
||||
x = upscale(self.keras, x, 256)
|
||||
|
||||
return self.keras.models.Model(input_layer, x)
|
||||
def Build(self, created_vram_gb):
|
||||
bgr_shape = (64, 64, 3)
|
||||
mask_shape = (64, 64, 1)
|
||||
|
||||
def Encoder(input_shape):
|
||||
input_layer = self.keras.layers.Input(input_shape)
|
||||
x = input_layer
|
||||
if created_vram_gb >= 4:
|
||||
x = conv(self.keras, x, 128)
|
||||
x = conv(self.keras, x, 256)
|
||||
x = conv(self.keras, x, 512)
|
||||
x = conv(self.keras, x, 1024)
|
||||
x = self.keras.layers.Dense(1024)(self.keras.layers.Flatten()(x))
|
||||
x = self.keras.layers.Dense(4 * 4 * 1024)(x)
|
||||
x = self.keras.layers.Reshape((4, 4, 1024))(x)
|
||||
x = upscale(self.keras, x, 512)
|
||||
else:
|
||||
x = conv(self.keras, x, 128 )
|
||||
x = conv(self.keras, x, 256 )
|
||||
x = conv(self.keras, x, 512 )
|
||||
x = conv(self.keras, x, 768 )
|
||||
x = self.keras.layers.Dense(512)(self.keras.layers.Flatten()(x))
|
||||
x = self.keras.layers.Dense(4 * 4 * 512)(x)
|
||||
x = self.keras.layers.Reshape((4, 4, 512))(x)
|
||||
x = upscale(self.keras, x, 256)
|
||||
|
||||
return self.keras.models.Model(input_layer, x)
|
||||
|
||||
def Decoder(self, created_vram_gb):
|
||||
if created_vram_gb >= 4:
|
||||
input_ = self.keras.layers.Input(shape=(8, 8, 512))
|
||||
else:
|
||||
input_ = self.keras.layers.Input(shape=(8, 8, 256))
|
||||
def Decoder():
|
||||
if created_vram_gb >= 4:
|
||||
input_ = self.keras.layers.Input(shape=(8, 8, 512))
|
||||
else:
|
||||
input_ = self.keras.layers.Input(shape=(8, 8, 256))
|
||||
|
||||
x = input_
|
||||
x = upscale(self.keras, x, 256)
|
||||
x = upscale(self.keras, x, 128)
|
||||
x = upscale(self.keras, x, 64)
|
||||
|
||||
x = input_
|
||||
x = upscale(self.keras, x, 256)
|
||||
x = upscale(self.keras, x, 128)
|
||||
x = upscale(self.keras, x, 64)
|
||||
|
||||
y = input_ #mask decoder
|
||||
y = upscale(self.keras, y, 256)
|
||||
y = upscale(self.keras, y, 128)
|
||||
y = upscale(self.keras, y, 64)
|
||||
|
||||
x = self.keras.layers.convolutional.Conv2D(3, kernel_size=5, padding='same', activation='sigmoid')(x)
|
||||
y = self.keras.layers.convolutional.Conv2D(1, kernel_size=5, padding='same', activation='sigmoid')(y)
|
||||
|
||||
|
||||
return self.keras.models.Model(input_, [x,y])
|
||||
y = input_ #mask decoder
|
||||
y = upscale(self.keras, y, 256)
|
||||
y = upscale(self.keras, y, 128)
|
||||
y = upscale(self.keras, y, 64)
|
||||
|
||||
x = self.keras.layers.convolutional.Conv2D(3, kernel_size=5, padding='same', activation='sigmoid')(x)
|
||||
y = self.keras.layers.convolutional.Conv2D(1, kernel_size=5, padding='same', activation='sigmoid')(y)
|
||||
|
||||
|
||||
return self.keras.models.Model(input_, [x,y])
|
||||
|
||||
return bgr_shape, mask_shape, Encoder(bgr_shape), Decoder(), Decoder()
|
||||
Loading…
Add table
Add a link
Reference in a new issue