def generator_model_factory(
input_params,
problem_type: pt.ProblemType,
):
if problem_type == pt.ProblemType.VANILLA_MNIST.name:
return latent_to_image.LatentToImageGenerator(input_params)
if problem_type == pt.ProblemType.VANILLA_FASHION_MNIST.name:
return latent_to_image.LatentToImageGenerator(input_params)
elif problem_type == pt.ProblemType.VANILLA_CIFAR10.name:
return latent_to_image.LatentToImageCifar10Generator(input_params)
elif problem_type == pt.ProblemType.CONDITIONAL_MNIST.name:
return conditional_latent_to_image.LatentToImageConditionalGenerator(
input_params)
elif problem_type == pt.ProblemType.CONDITIONAL_FASHION_MNIST.name:
return conditional_latent_to_image.LatentToImageConditionalGenerator(
input_params)
elif problem_type == pt.ProblemType.CONDITIONAL_CIFAR10.name:
return conditional_latent_to_image.LatentToImageCifar10CConditionalGenerator(
input_params)
elif problem_type == pt.ProblemType.CYCLE_SUMMER2WINTER.name:
return [
unet.UNetGenerator(input_params),
unet.UNetGenerator(input_params)
]
else:
raise NotImplementedError
def test_random_to_image_generator_output_shape(self):
input_params = edict({'latent_size': 100})
g = latent_to_image.LatentToImageGenerator(input_params)
z = tf.random.normal(shape=[1, 100])
output_img = g(z)
expected_shape = (1, 28, 28, 1)
self.assertEqual(output_img.shape, expected_shape)
'save_images_every_n_steps': 10
})
dataset = mnist.MnistDataset(model_parameters, with_labels=True)
def validation_dataset():
test_batch_size = model_parameters.num_classes ** 2
labels = np.repeat(list(range(model_parameters.num_classes)), model_parameters.num_classes)
validation_samples = [tf.random.normal([test_batch_size, model_parameters.latent_size]), np.array(labels)]
return validation_samples
validation_dataset = validation_dataset()
generator = latent_to_image.LatentToImageGenerator(model_parameters)
discriminator = discriminator.Discriminator(model_parameters)
generator_optimizer = optimizers.Adam(
learning_rate=model_parameters.learning_rate_generator,
beta_1=0.5,
)
discriminator_optimizer = optimizers.Adam(
learning_rate=model_parameters.learning_rate_discriminator,
beta_1=0.5,
)
callbacks = [
saver.ImageProblemSaver(
save_images_every_n_steps=model_parameters.save_images_every_n_steps,
)