def func(input):
b, h, w, c = K.int_shape(input)
n_downs = get_power_of_two(w) - 4
Norm = XNormalizationL()
Norm2 = XNormalizationL()
Norm4 = XNormalizationL()
Norm8 = XNormalizationL()
x = input
x = e1 = XConv2D(ngf, 4, strides=2, use_bias=True)(x)
x = e2 = Norm2(XConv2D(ngf * 2, 4, strides=2)(LeakyReLU(0.2)(x)))
x = e3 = Norm4(XConv2D(ngf * 4, 4, strides=2)(LeakyReLU(0.2)(x)))
l = []
for i in range(n_downs):
x = Norm8(XConv2D(ngf * 8, 4, strides=2)(LeakyReLU(0.2)(x)))
l += [x]
x = XConv2D(ngf * 8, 4, strides=2,
use_bias=True)(LeakyReLU(0.2)(x))
for i in range(n_downs):
x = Norm8(XConv2DTranspose(ngf * 8, 4, strides=2)(ReLU()(x)))
if i <= n_downs - 2:
x = Dropout(0.5)(x)
x = Concatenate(axis=-1)([x, l[-i - 1]])
x = Norm4(XConv2DTranspose(ngf * 4, 4, strides=2)(ReLU()(x)))
x = Concatenate(axis=-1)([x, e3])
x = Norm2(XConv2DTranspose(ngf * 2, 4, strides=2)(ReLU()(x)))
x = Concatenate(axis=-1)([x, e2])
x = Norm(XConv2DTranspose(ngf, 4, strides=2)(ReLU()(x)))
x = Concatenate(axis=-1)([x, e1])
x = XConv2DTranspose(output_nc,
4,
strides=2,
activation='tanh',
use_bias=True)(ReLU()(x))
return x
def onInitialize(self, batch_size=-1, **in_options):
exec(nnlib.code_import_all, locals(), globals())
created_batch_size = self.get_batch_size()
if self.epoch == 0:
#first run
print("\nModel first run. Enter options.")
try:
created_resolution = int(
input("Resolution (default:64, valid: 64,128,256) : "))
except:
created_resolution = 64
if created_resolution not in [64, 128, 256]:
created_resolution = 64
try:
created_batch_size = int(
input("Batch_size (minimum/default - 10) : "))
except:
created_batch_size = 10
created_batch_size = max(created_batch_size, 1)
print("Done. If training won't start, decrease resolution")
self.options['created_resolution'] = created_resolution
self.options['created_batch_size'] = created_batch_size
self.created_vram_gb = self.device_config.gpu_total_vram_gb
else:
#not first run
if 'created_batch_size' in self.options.keys():
created_batch_size = self.options['created_batch_size']
else:
raise Exception(
"Continue training, but created_batch_size not found.")
if 'created_resolution' in self.options.keys():
created_resolution = self.options['created_resolution']
else:
raise Exception(
"Continue training, but created_resolution not found.")
resolution = created_resolution
bgr_shape = (resolution, resolution, 3)
ngf = 64
npf = 64
ndf = 64
lambda_A = 10
lambda_B = 10
self.set_batch_size(created_batch_size)
use_batch_norm = created_batch_size > 1
self.GA = modelify(
ResNet(bgr_shape[2],
use_batch_norm,
n_blocks=6,
ngf=ngf,
use_dropout=False))(Input(bgr_shape))
self.GB = modelify(
ResNet(bgr_shape[2],
use_batch_norm,
n_blocks=6,
ngf=ngf,
use_dropout=False))(Input(bgr_shape))
#self.GA = modelify(UNet (bgr_shape[2], use_batch_norm, num_downs=get_power_of_two(resolution)-1, ngf=ngf, use_dropout=True))(Input(bgr_shape))
#self.GB = modelify(UNet (bgr_shape[2], use_batch_norm, num_downs=get_power_of_two(resolution)-1, ngf=ngf, use_dropout=True))(Input(bgr_shape))
self.PA = modelify(
UNetTemporalPredictor(
bgr_shape[2],
use_batch_norm,
num_downs=get_power_of_two(resolution) - 1,
ngf=npf,
use_dropout=True))([Input(bgr_shape),
Input(bgr_shape)])
self.PB = modelify(
UNetTemporalPredictor(
bgr_shape[2],
use_batch_norm,
num_downs=get_power_of_two(resolution) - 1,
ngf=npf,
use_dropout=True))([Input(bgr_shape),
Input(bgr_shape)])
self.DA = modelify(
NLayerDiscriminator(use_batch_norm, ndf=ndf,
n_layers=3))(Input(bgr_shape))
self.DB = modelify(
NLayerDiscriminator(use_batch_norm, ndf=ndf,
n_layers=3))(Input(bgr_shape))
if not self.is_first_run():
self.GA.load_weights(self.get_strpath_storage_for_file(self.GAH5))
self.DA.load_weights(self.get_strpath_storage_for_file(self.DAH5))
self.PA.load_weights(self.get_strpath_storage_for_file(self.PAH5))
self.GB.load_weights(self.get_strpath_storage_for_file(self.GBH5))
self.DB.load_weights(self.get_strpath_storage_for_file(self.DBH5))
self.PB.load_weights(self.get_strpath_storage_for_file(self.PBH5))
real_A0 = Input(bgr_shape, name="real_A0")
real_A1 = Input(bgr_shape, name="real_A1")
real_A2 = Input(bgr_shape, name="real_A2")
real_B0 = Input(bgr_shape, name="real_B0")
real_B1 = Input(bgr_shape, name="real_B1")
real_B2 = Input(bgr_shape, name="real_B2")
DA_ones = K.ones(K.int_shape(self.DA.outputs[0])[1:])
DA_zeros = K.zeros(K.int_shape(self.DA.outputs[0])[1:])
DB_ones = K.ones(K.int_shape(self.DB.outputs[0])[1:])
DB_zeros = K.zeros(K.int_shape(self.DB.outputs[0])[1:])
def CycleLoss(t1, t2):
return K.mean(K.square(t1 - t2))
def RecurrentLOSS(t1, t2):
return K.mean(K.square(t1 - t2))
def RecycleLOSS(t1, t2):
return K.mean(K.square(t1 - t2))
fake_B0 = self.GA(real_A0)
fake_B1 = self.GA(real_A1)
fake_A0 = self.GB(real_B0)
fake_A1 = self.GB(real_B1)
#rec_FB0 = self.GA(fake_A0)
#rec_FB1 = self.GA(fake_A1)
#rec_FA0 = self.GB(fake_B0)
#rec_FA1 = self.GB(fake_B1)
pred_A2 = self.PA([real_A0, real_A1])
pred_B2 = self.PB([real_B0, real_B1])
rec_A2 = self.GB(self.PB([fake_B0, fake_B1]))
rec_B2 = self.GA(self.PA([fake_A0, fake_A1]))
loss_G = K.mean(K.square(self.DB(fake_B0) - DB_ones)) + \
K.mean(K.square(self.DB(fake_B1) - DB_ones)) + \
K.mean(K.square(self.DA(fake_A0) - DA_ones)) + \
K.mean(K.square(self.DA(fake_A1) - DA_ones)) + \
lambda_A * ( #CycleLoss(rec_FA0, real_A0) + \
#CycleLoss(rec_FA1, real_A1) + \
RecurrentLOSS(pred_A2, real_A2) + \
RecycleLOSS(rec_A2, real_A2) ) + \
lambda_B * ( #CycleLoss(rec_FB0, real_B0) + \
#CycleLoss(rec_FB1, real_B1) + \
RecurrentLOSS(pred_B2, real_B2) + \
RecycleLOSS(rec_B2, real_B2) )
weights_G = self.GA.trainable_weights + self.GB.trainable_weights + self.PA.trainable_weights + self.PB.trainable_weights
self.G_train = K.function(
[real_A0, real_A1, real_A2, real_B0, real_B1, real_B2], [loss_G],
Adam(lr=2e-4, beta_1=0.5,
beta_2=0.999).get_updates(loss_G, weights_G))
###########
loss_D_A0 = ( K.mean(K.square( self.DA(real_A0) - DA_ones)) + \
K.mean(K.square( self.DA(fake_A0) - DA_zeros)) ) * 0.5
loss_D_A1 = ( K.mean(K.square( self.DA(real_A1) - DA_ones)) + \
K.mean(K.square( self.DA(fake_A1) - DA_zeros)) ) * 0.5
loss_D_A = loss_D_A0 + loss_D_A1
self.DA_train = K.function(
[real_A0, real_A1, real_A2, real_B0, real_B1, real_B2], [loss_D_A],
Adam(lr=2e-4, beta_1=0.5,
beta_2=0.999).get_updates(loss_D_A,
self.DA.trainable_weights))
############
loss_D_B0 = ( K.mean(K.square( self.DB(real_B0) - DB_ones)) + \
K.mean(K.square( self.DB(fake_B0) - DB_zeros)) ) * 0.5
loss_D_B1 = ( K.mean(K.square( self.DB(real_B1) - DB_ones)) + \
K.mean(K.square( self.DB(fake_B1) - DB_zeros)) ) * 0.5
loss_D_B = loss_D_B0 + loss_D_B1
self.DB_train = K.function(
[real_A0, real_A1, real_A2, real_B0, real_B1, real_B2], [loss_D_B],
Adam(lr=2e-4, beta_1=0.5,
beta_2=0.999).get_updates(loss_D_B,
self.DB.trainable_weights))
############
self.G_view = K.function(
[real_A0, real_A1, real_A2, real_B0, real_B1, real_B2], [
fake_A0, fake_A1, pred_A2, rec_A2, fake_B0, fake_B1, pred_B2,
rec_B2
])
self.G_convert = K.function([real_B0], [fake_A0])
if self.is_training_mode:
f = SampleProcessor.TypeFlags
self.set_training_data_generators([
SampleGeneratorImageTemporal(
self.training_data_src_path,
debug=self.is_debug(),
batch_size=self.batch_size,
temporal_image_count=3,
sample_process_options=SampleProcessor.Options(
random_flip=False, normalize_tanh=True),
output_sample_types=[[f.SOURCE | f.MODE_BGR, resolution]]),
SampleGeneratorImageTemporal(
self.training_data_dst_path,
debug=self.is_debug(),
batch_size=self.batch_size,
temporal_image_count=3,
sample_process_options=SampleProcessor.Options(
random_flip=False, normalize_tanh=True),
output_sample_types=[[f.SOURCE | f.MODE_BGR, resolution]]),
])
def onInitialize(self, batch_size=-1, **in_options):
exec(nnlib.code_import_all, locals(), globals())
self.set_vram_batch_requirements({2: 64})
resolution = self.options['resolution']
bgr_shape = (resolution, resolution, 3)
mask_shape = (resolution, resolution, 1)
ngf = 64
npf = 64
ndf = 64
lambda_A = 10
lambda_B = 10
#self.set_batch_size(created_batch_size)
#self.GA = modelify(ResNet (bgr_shape[2], use_batch_norm, n_blocks=6, ngf=ngf, use_dropout=True))(Input(bgr_shape))
#self.GB = modelify(ResNet (bgr_shape[2], use_batch_norm, n_blocks=6, ngf=ngf, use_dropout=True))(Input(bgr_shape))
#self.GA = modelify(TestModel.GFlow (bgr_shape[2], use_batch_norm, num_downs=get_power_of_two(resolution)-1, ngf=ngf, use_dropout=True))(Input(bgr_shape))
#self.GB = modelify(UNet (bgr_shape[2], use_batch_norm, num_downs=get_power_of_two(resolution)-1, ngf=ngf, use_dropout=True))(Input(bgr_shape))
self.G = modelify(
TestModel.GFlow(1,
True,
num_downs=get_power_of_two(resolution) - 1,
ngf=64))(Input(bgr_shape))
#self.enc = modelify ( TestModel.EncFlow(resolution, 3, ae_dims=256, ed_ch_dims=21 ) ) (Input(bgr_shape))
#dec_Inputs = [ Input(K.int_shape(x)[1:]) for x in self.enc.outputs ]
#self.dec = modelify ( TestModel.DecFlow(resolution, 1, ae_dims=256, ed_ch_dims=21 ) ) (dec_Inputs)
#self.DA = modelify(NLayerDiscriminator(use_batch_norm, ndf=ndf, n_layers=3) ) (Input(bgr_shape))
#self.DB = modelify(NLayerDiscriminator(use_batch_norm, ndf=ndf, n_layers=3) ) (Input(bgr_shape))
if not self.is_first_run():
self.G.load_weights(self.get_strpath_storage_for_file(self.GH5))
#self.enc.load_weights (self.get_strpath_storage_for_file(self.encH5))
#self.dec.load_weights (self.get_strpath_storage_for_file(self.decH5))
#self.GA.load_weights (self.get_strpath_storage_for_file(self.GAH5))
#self.GB.load_weights (self.get_strpath_storage_for_file(self.GBH5))
#self.DA.load_weights (self.get_strpath_storage_for_file(self.DAH5))
#self.DB.load_weights (self.get_strpath_storage_for_file(self.DBH5))
#import code
#code.interact(local=dict(globals(), **locals()))
warped_A0 = Input(bgr_shape)
real_A0 = Input(bgr_shape)
mask_A0 = Input(mask_shape)
#code_A0 = self.enc(warped_A0)
#rec_A0 = self.dec(code_A0)
rec_A0 = self.G(warped_A0)
loss_G = K.mean(K.square(mask_A0 - rec_A0))
if self.is_training_mode:
def optimizer():
return Adam(lr=5e-5, beta_1=0.5, beta_2=0.999)
#self.GA_train = K.function ([real_A0, real_B0],[loss_GA], optimizer().get_updates(loss_GA, self.GA.trainable_weights + self.GB.trainable_weights) )
#self.GB_train = K.function ([real_A0, real_B0],[loss_GB], optimizer().get_updates(loss_GB, self.GA.trainable_weights + self.GB.trainable_weights) )
self.G_train = K.function([warped_A0, mask_A0], [loss_G],
optimizer().get_updates(
loss_G, self.G.trainable_weights))
#self.dec_train = K.function ([warped_A0, real_A0, mask_A0, real_pitch_A0, real_yaw_A0],[loss_G], optimizer().get_updates(loss_G, self.dec.trainable_weights ) )
#self.GB_train = K.function ([real_A0, real_B0],[loss_GB], optimizer().get_updates(loss_GB, self.G.trainable_weights) )
############
#
#loss_D_A = ( K.mean(K.square( self.DA(real_A0) )) + \
# K.mean(K.square( self.DA(fake_A0) - DA_ones)) ) * 0.5
#
#self.DA_train = K.function ([real_A0, real_B0],[loss_D_A],
# optimizer().get_updates(loss_D_A, self.DA.trainable_weights) )
#
#############
#
#loss_D_B = ( K.mean(K.square( self.DB(real_B0) )) + \
# K.mean(K.square( self.DB(fake_B0) - DB_ones)) ) * 0.5
#
#self.DB_train = K.function ([real_A0, real_B0],[loss_D_B],
# optimizer().get_updates(loss_D_B, self.DB.trainable_weights) )
#
############
self.G_view = K.function([warped_A0], [rec_A0])
else:
self.G_convert = K.function([real_B0], [fake_A0])
if self.is_training_mode:
f = SampleProcessor.TypeFlags
face_type = f.FACE_ALIGN_FULL if self.options[
'face_type'] == 'f' else f.FACE_ALIGN_HALF
self.set_training_data_generators([
SampleGeneratorFace(
self.training_data_src_path,
sort_by_yaw_target_samples_path=self.training_data_dst_path
if self.sort_by_yaw else None,
debug=self.is_debug(),
batch_size=self.batch_size,
sample_process_options=SampleProcessor.Options(
random_flip=self.random_flip,
normalize_tanh=True,
scale_range=np.array([-0.05, 0.05]) +
self.src_scale_mod / 100.0),
output_sample_types=[[
f.WARPED_TRANSFORMED | face_type | f.MODE_BGR,
resolution
], [f.TRANSFORMED | face_type | f.MODE_BGR, resolution],
[
f.TRANSFORMED | face_type
| f.MODE_M | f.FACE_MASK_FULL,
resolution
]],
add_pitch=True,
add_yaw=True),
SampleGeneratorFace(
self.training_data_dst_path,
debug=self.is_debug(),
batch_size=self.batch_size,
sample_process_options=SampleProcessor.Options(
random_flip=self.random_flip, normalize_tanh=True),
output_sample_types=[[
f.WARPED_TRANSFORMED | face_type | f.MODE_BGR,
resolution
], [f.TRANSFORMED | face_type | f.MODE_BGR, resolution],
[
f.TRANSFORMED | face_type
| f.MODE_M | f.FACE_MASK_FULL,
resolution
]],
add_pitch=True,
add_yaw=True)
])
