def test_kernel_shape_correct(self):
conv2d_layer = keras_layers.FanInScaledConv2D(filters=5,
kernel_size=(2, 3))
conv2d_layer(tf.ones((1, 8, 6, 4)))
self.evaluate(tf.compat.v1.global_variables_initializer())
self.assertSequenceEqual(conv2d_layer.kernel.shape, (2, 3, 4, 5))
def test_bias_has_correct_value(self, bias_multiplier=2.0):
conv2d_layer = keras_layers.FanInScaledConv2D(
filters=1,
kernel_size=3,
bias_initializer='ones',
bias_multiplier=bias_multiplier)
conv2d_layer(tf.ones((1, 3, 3, 4)))
self.evaluate(tf.compat.v1.global_variables_initializer())
self.assertAllClose(tf.squeeze(conv2d_layer.bias), bias_multiplier)
def test_kernel_has_correct_value(self, kernel_multiplier):
conv2d_layer = keras_layers.FanInScaledConv2D(
filters=5,
kernel_size=(2, 3),
kernel_initializer='ones',
kernel_multiplier=kernel_multiplier)
conv2d_layer(tf.ones((1, 8, 6, 4)))
self.evaluate(tf.compat.v1.global_variables_initializer())
# Checking if the 1 is correctly multiplied with sqrt(2/fan_in).
self.assertAllClose(
conv2d_layer.kernel,
tf.ones(shape=(2, 3, 4, 5)) * math.sqrt(2.0 / (2.0 * 3.0 * 4.0)) *
kernel_multiplier)
def create_conv_layer(use_fan_in_scaled_kernel: bool = False,
multiplier: float = math.sqrt(2),
**kwargs) -> tf.keras.layers.Conv2D:
"""Creates a convolutional layer.
Args:
use_fan_in_scaled_kernel: Whether to use a FanInScaledConv2D or a standard
Conv2D layer.
multiplier: Additional multiplier used only for FanInSclaedConv2D layer.
**kwargs: Keyword arguments forwarded to the convolutional layers.
Returns:
The created convolutional layer instance.
"""
if use_fan_in_scaled_kernel:
return keras_layers.FanInScaledConv2D(multiplier=multiplier, **kwargs)
else:
return tf.keras.layers.Conv2D(**kwargs)
def test_config_contains_both_base_classes(self):
conv2d_layer = keras_layers.FanInScaledConv2D(filters=3,
kernel_size=1,
kernel_multiplier=2.0,
multiplier=1.0)
config_dict = conv2d_layer.get_config()
with self.subTest(name='_FanInScaler'):
self.assertIn('kernel_multiplier', config_dict)
self.assertIn('multiplier', config_dict)
self.assertIn('bias_multiplier', config_dict)
self.assertEqual(config_dict['kernel_multiplier'], 2.0)
self.assertEqual(config_dict['multiplier'], 1.0)
self.assertIsNone(config_dict['bias_multiplier'])
with self.subTest(name='Conv2D'):
self.assertIn('filters', config_dict)
self.assertEqual(config_dict['filters'], 3)
def to_rgb(input_tensor: tf.Tensor,
kernel_initializer: _KerasInitializer,
name: Optional[str] = None) -> tf.Tensor:
"""Converts a feature map to an rgb output.
Args:
input_tensor: The input feature map.
kernel_initializer: The kernel initializer to use.
name: The name of the layer.
Returns:
The rgb image.
"""
return keras_layers.FanInScaledConv2D(
multiplier=1.0,
filters=3,
kernel_size=1,
strides=1,
kernel_initializer=kernel_initializer,
padding='same',
name=name)(input_tensor)
def create_generator(
latent_code_dimension: int = 128,
upsampling_blocks_num_channels: Sequence[int] = (512, 256, 128, 64),
relu_leakiness: float = 0.2,
kernel_initializer: Optional[_KerasInitializer] = None,
use_pixel_normalization: bool = True,
use_batch_normalization: bool = False,
generate_intermediate_outputs: bool = False,
normalize_latent_code: bool = True,
name: str = 'progressive_gan_generator') -> tf.keras.Model:
"""Creates a Keras model for the generator network architecture.
This architecture is implemented according to the paper "Progressive growing
of GANs for Improved Quality, Stability, and Variation"
https://arxiv.org/abs/1710.10196
The intermediate outputs are optionally provided for the architecture of
"MSG-GAN: Multi-Scale Gradient GAN for Stable Image Synthesis"
https://arxiv.org/abs/1903.06048
Args:
latent_code_dimension: The number of dimensions in the latent code.
upsampling_blocks_num_channels: The number of channels for each upsampling
block. This argument also determines how many upsampling blocks are added.
relu_leakiness: Slope of the negative part of the leaky relu.
kernel_initializer: Initializer of the kernel. If none TruncatedNormal is
used.
use_pixel_normalization: If pixel normalization layers should be inserted to
the network.
use_batch_normalization: If batch normalization layers should be inserted to
the network.
generate_intermediate_outputs: If true the model outputs a list of
tf.Tensors with increasing resolution starting with the starting_size up
to the final resolution output.
normalize_latent_code: If true the latent code is normalized to unit length
before feeding it to the network.
name: The name of the Keras model.
Returns:
The created generator keras model object.
"""
if kernel_initializer is None:
kernel_initializer = tf.keras.initializers.TruncatedNormal(mean=0.0,
stddev=1.0)
input_tensor = tf.keras.Input(shape=(latent_code_dimension, ))
if normalize_latent_code:
maybe_normzlized_input_tensor = keras_layers.PixelNormalization(
axis=1)(input_tensor)
else:
maybe_normzlized_input_tensor = input_tensor
tensor = keras_layers.FanInScaledDense(
multiplier=math.sqrt(2.0) / 4.0,
units=4 * 4 * latent_code_dimension,
kernel_initializer=kernel_initializer)(maybe_normzlized_input_tensor)
tensor = tf.keras.layers.Reshape(
target_shape=(4, 4, latent_code_dimension))(tensor)
tensor = tf.keras.layers.LeakyReLU(alpha=relu_leakiness)(tensor)
if use_batch_normalization:
tensor = tf.keras.layers.BatchNormalization()(tensor)
if use_pixel_normalization:
tensor = keras_layers.PixelNormalization(axis=3)(tensor)
tensor = keras_layers.FanInScaledConv2D(
filters=upsampling_blocks_num_channels[0],
kernel_size=3,
strides=1,
padding='same',
kernel_initializer=kernel_initializer)(tensor)
tensor = tf.keras.layers.LeakyReLU(alpha=relu_leakiness)(tensor)
if use_batch_normalization:
tensor = tf.keras.layers.BatchNormalization()(tensor)
if use_pixel_normalization:
tensor = keras_layers.PixelNormalization(axis=3)(tensor)
outputs = []
for index, channels in enumerate(upsampling_blocks_num_channels):
if generate_intermediate_outputs:
outputs.append(
to_rgb(input_tensor=tensor,
kernel_initializer=kernel_initializer,
name='side_output_%d_conv' % index))
tensor = keras_layers.TwoByTwoNearestNeighborUpSampling()(tensor)
for _ in range(2):
tensor = keras_layers.FanInScaledConv2D(
filters=channels,
kernel_size=3,
strides=1,
padding='same',
kernel_initializer=kernel_initializer)(tensor)
tensor = tf.keras.layers.LeakyReLU(alpha=relu_leakiness)(tensor)
if use_batch_normalization:
tensor = tf.keras.layers.BatchNormalization()(tensor)
if use_pixel_normalization:
tensor = keras_layers.PixelNormalization(axis=3)(tensor)
tensor = to_rgb(input_tensor=tensor,
kernel_initializer=kernel_initializer,
name='final_output')
if generate_intermediate_outputs:
outputs.append(tensor)
return tf.keras.Model(inputs=input_tensor, outputs=outputs, name=name)
else:
return tf.keras.Model(inputs=input_tensor, outputs=tensor, name=name)
def create_synthesis_network(latent_code_dimension: int = 128,
upsampling_blocks_num_channels: Sequence[int] = (
512, 256, 128, 64),
relu_leakiness: float = 0.2,
generate_intermediate_outputs: bool = False,
use_bilinear_upsampling: bool = True,
name: str = 'synthesis') -> tf.keras.Model:
"""Creates the synthesis network using the functional API.
The function creates the synthesis network as defined in "A Style-Based
Generator Architecture for Generative Adversarial Networks"
https://arxiv.org/abs/1812.04948 using the Keras functional API.
Args:
latent_code_dimension: The number of dimensions in the latent code.
upsampling_blocks_num_channels: The number of channels for each upsampling
block. This argument also determines how many upsampling blocks are added.
relu_leakiness: Slope of the negative part of the leaky relu.
generate_intermediate_outputs: If true the model outputs a list of
tf.Tensors with increasing resolution starting with the starting_size up
to the final resolution output.
use_bilinear_upsampling: If true bilinear upsampling is used.
name: The name of the Keras model.
Returns:
The synthesis network.
"""
kernel_initializer = tf.keras.initializers.TruncatedNormal(
mean=0.0, stddev=1.0)
mapped_latent_code_input = tf.keras.Input(shape=(latent_code_dimension,))
tensor = keras_layers.LearnedConstant()(mapped_latent_code_input)
tensor = keras_layers.Noise()(tensor)
tensor = tf.keras.layers.LeakyReLU(alpha=relu_leakiness)(tensor)
tensor = apply_style_with_adain(
mapped_latent_code=mapped_latent_code_input, input_tensor=tensor)
tensor = keras_layers.FanInScaledConv2D(
filters=upsampling_blocks_num_channels[0],
kernel_size=3,
strides=1,
padding='same',
kernel_initializer=kernel_initializer)(
tensor)
tensor = keras_layers.Noise()(tensor)
tensor = tf.keras.layers.LeakyReLU(alpha=relu_leakiness)(tensor)
tensor = apply_style_with_adain(
mapped_latent_code=mapped_latent_code_input, input_tensor=tensor)
outputs = []
for index, channels in enumerate(upsampling_blocks_num_channels):
if generate_intermediate_outputs:
outputs.append(
architectures_progressive_gan.to_rgb(
input_tensor=tensor,
kernel_initializer=kernel_initializer,
name='side_output_%d_conv' % index))
tensor = keras_layers.TwoByTwoNearestNeighborUpSampling()(tensor)
if use_bilinear_upsampling:
tensor = keras_layers.Blur2D()(tensor)
for _ in range(2):
tensor = keras_layers.FanInScaledConv2D(
filters=channels,
kernel_size=3,
strides=1,
padding='same',
kernel_initializer=kernel_initializer)(
tensor)
tensor = keras_layers.Noise()(tensor)
tensor = tf.keras.layers.LeakyReLU(alpha=relu_leakiness)(tensor)
tensor = apply_style_with_adain(
mapped_latent_code=mapped_latent_code_input, input_tensor=tensor)
tensor = architectures_progressive_gan.to_rgb(
input_tensor=tensor,
kernel_initializer=kernel_initializer,
name='final_output')
if generate_intermediate_outputs:
outputs.append(tensor)
return tf.keras.Model(
inputs=mapped_latent_code_input, outputs=outputs, name=name)
else:
return tf.keras.Model(
inputs=mapped_latent_code_input, outputs=tensor, name=name)