def build_discriminator_model(discriminator: Model,
resolution: int,
discriminator_lr: float,
loss_exponent: int,
channels: int = 3) -> Model:
set_model_trainable(discriminator, True)
real_samples = Input((resolution, resolution, channels))
generated_samples = Input((resolution, resolution, channels))
k = Input((1, ))
discriminated_real_samples = discriminator(real_samples)
discriminated_generated_samples = discriminator(generated_samples)
ln_real = ln_loss(real_samples, discriminated_real_samples, loss_exponent)
ln_generated = ln_loss(generated_samples, discriminated_generated_samples,
loss_exponent)
discriminator_model = Model([real_samples, generated_samples, k],
[ln_real, ln_generated],
name='discriminator_model')
discriminator_model.compile(optimizer=Adam(discriminator_lr),
loss=discriminator_loss(
k, ln_real, ln_generated))
return discriminator_model
def build_generator_model(generator: Model, discriminator: Model,
latent_dim: int, generator_lr: float,
loss_exponent: int) -> Model:
set_model_trainable(discriminator, False)
input_noise = Input((latent_dim, ))
generated_samples = generator(input_noise)
discriminated_generated_samples = discriminator(generated_samples)
ln_generated = ln_loss(generated_samples, discriminated_generated_samples,
loss_exponent)
generator_model = Model(input_noise, ln_generated, name='generator_model')
generator_model.compile(optimizer=Adam(generator_lr),
loss=generator_loss(ln_generated))
return generator_model
def build_generator_model(generator: Model, critic: Model, latent_dim: int,
generator_lr: float) -> Model:
set_model_trainable(generator, True)
set_model_trainable(critic, False)
noise_samples = Input((latent_dim, ))
generated_samples = generator(noise_samples)
generated_criticized = critic(generated_samples)
generator_model = Model(noise_samples,
generated_criticized,
name='generator_model')
generator_model.compile(optimizer=Adam(generator_lr, beta_1=0, beta_2=0.9),
loss=wasserstein_loss)
return generator_model
def build_encoder_decoder_models(encoder: Model, decoder_generator: Model,
critic: Model, resolution: int,
latent_dim: int, channels: int, gamma: float,
vae_lr: float) -> Tuple[Model, Model]:
set_model_trainable(encoder, True)
set_model_trainable(decoder_generator, True)
set_model_trainable(critic, False)
real_samples = Input((resolution, resolution, channels))
noise_samples = Input((latent_dim, ))
generated_samples = decoder_generator(noise_samples)
generated_criticized = critic(generated_samples)
real_criticized = critic(real_samples)
z_mean, z_log_var = encoder(real_samples)
sampled_z = Lambda(sampling)([z_mean, z_log_var])
decoded_samples = decoder_generator(sampled_z)
decoded_criticized = critic(decoded_samples)
vae_model = Model(
[real_samples, noise_samples],
[generated_criticized, generated_criticized, generated_criticized])
vae_model.compile(
optimizer=Adam(lr=vae_lr, beta_1=0.5, beta_2=0.9),
loss=[
kl_loss(z_mean, z_log_var), wasserstein_loss,
mse_loss(real_criticized, decoded_criticized)
],
loss_weights=[1 / 3.0, gamma * 1 / 3.0, (1 - gamma) * 1 / 3.0])
generator = Model(noise_samples, generated_samples)
return vae_model, generator
def build_critic_model(generator: Model,
critic: Model,
latent_dim: int,
resolution: int,
classes_n: int,
batch_size: int,
critic_lr: float,
gradient_penalty_weight: int,
channels: int = 3) -> Model:
set_model_trainable(generator, False)
set_model_trainable(critic, True)
noise_samples = Input((latent_dim, ))
class_samples = Input((classes_n, ))
real_samples = Input((resolution, resolution, channels))
generated_samples = generator([noise_samples, class_samples])
generated_criticized = critic([generated_samples, class_samples])
real_criticized = critic([real_samples, class_samples])
averaged_samples = RandomWeightedAverage(batch_size)(
[real_samples, generated_samples])
averaged_criticized = critic([averaged_samples, class_samples])
partial_gp_loss = partial(gradient_penalty_loss,
averaged_samples=averaged_samples,
gradient_penalty_weight=gradient_penalty_weight)
partial_gp_loss.__name__ = 'gradient_penalty'
critic_model = Model(
[real_samples, noise_samples, class_samples],
[real_criticized, generated_criticized, averaged_criticized],
name='critic_model')
critic_model.compile(
optimizer=Adam(critic_lr, beta_1=0.5, beta_2=0.9),
loss=[wasserstein_loss, wasserstein_loss, partial_gp_loss])
return critic_model