def info_recognition(input_tensor,
cat_code_dim=0,
cont_code_dim=0,
dataset='mnist'):
"""Recognition network Q as proposed in section 3 of "InfoGAN:
Interpretable Representation Learning by Information Maximizing Generative
Adversarial Nets" (https://arxiv.org/pdf/1606.03657.pdf).
"""
x = _conv2d_bn(input_tensor,
64,
4,
strides=2,
bn_before=None,
activation='leaky_relu')
x = _conv2d_bn(x,
128,
4,
strides=2,
bn_before=True,
activation='leaky_relu')
x = _dense(x, 1024, bn_before=True, activation='leaky_relu')
x = _dense(x, 128, bn_before=True, activation='leaky_relu')
x = _dense(x, 1 + cat_code_dim + cont_code_dim, activation='linear')
return x
def auto_decode(model, latent_tensor, output_shape):
"""Automaticly builds the core layers of a decoder given the encoder.
The list returned stores a dictionary of layer configs to be passed in to
transposed conv layers. Any max pooling layers in the encoder will be
replaced with deconv layers as well.
Parameters
----------
model : tf.keras.Model
Encoder to tranpose.
latent_tensor : tf.Tensor
Latent tensor of shape (batch_size, latent_dim) used as input to
the decoder.
output_shape : tuple
Shape of the input to reconstruct.
Returns
-------
x : tf.Tensor
Output of decoder with the same shape as the input for the encoder.
"""
conv_configs, conv_shapes = _conv_config_shape(model.layers)
final_conv_shape = conv_shapes[0]
x = _dense(latent_tensor, 256)
x = _dense(x, np.prod(final_conv_shape))
x = tf.reshape(x, [-1] + final_conv_shape)
for config in conv_configs:
x = _deconv2d_bn(x, **config)
return x
def info_generator(latent_tensor, dataset='mnist', name='gen_'):
"""Builds a infogan generator."""
x = tf.reshape(latent_tensor, shape=(-1, np.prod(latent_tensor.shape[1:])))
x = _dense(x, 1024, bn_before=True, activation='relu')
x = _dense(x, 128 * 7 * 7, bn_before=True, activation='relu')
x = _deconv2d_bn(x, 64, 4, strides=2, activation='relu')
x = _deconv2d_bn(x, 1, 4, strides=2, activation=None)
x = tf.keras.layers.Flatten()(x)
return x
def build_inference_model(input_tensor,
encoder_fn,
latent_dim,
name='encoder',
**kwargs):
"""Builds a Gaussian inference model (encoder), which takes in an
`input_tensor`, applies an `encoder_fn`, and returns the posterior q(z|x),
which is parameterized using a mean and log variance FC layer.
Parameters
----------
input_tensor : tf.Tensor
Input tensor to encode.
encoder_fn : function
Encoder function that expects an input tensor `x` and returns
the output tensor of encoder that will be connected to a mean and
log variance dense layers parameterizing the latent space. Example:
input_tensor = tf.keras.layers.Input((64, 64, 3))
x = encoder_fn(input_tensor)
z_mean = tf.keras.layers.Dense(latent_dim)(x)
z_log_var = tf.keras.layers.Dense(latent_dim)(x)
latent_dim : int
Dimension of latent space.
name : str (default='encoder')
Name of inference model.
Returns
-------
model : tf.keras.Model
Gaussian encoder which accepts `input_tensor` and outputs two tensors:
z_mean : Mean latent vector of the encoder.
z_log_var : Log variance latent vector of the encoder.
"""
x = encoder_fn(input_tensor, **kwargs)
if len(x.shape) > 2:
x = tf.keras.layers.Flatten()(x)
# parameterize the latent space with a mean a log variance FC layer.
z_mean = _dense(x,
latent_dim,
activation=None,
use_bias=False,
name='enc_mean')
z_log_var = _dense(x,
latent_dim,
activation=None,
use_bias=False,
name='enc_logvar')
model = tf.keras.models.Model(input_tensor, [z_mean, z_log_var], name=name)
return model
def fc_decoder(latent_tensor,
output_shape,
hidden_units=1200,
activation='tanh',
use_bias=False,
layer_prefix='dec_'):
"""MLP decoder as proposed in section A.4 of "beta-VAE: Learning Basic
Visual Concepts with a Constrained Variational Framework"
(https://openreview.net/references/pdf?id=Sy2fzU9gl).
Parameters
----------
latent_tensor : tf.Tensor
Input tensor to encode.
output_shape : tuple
Shape of input to reconstruct.
hidden_units : int (default=1200)
Number of hidden units in each FC layer.
activation : str, tf.keras.layers.Activation, tf.nn (default='tanh')
Activation function applied to linear dense layer.
use_bias : bool (default=False)
Boolean whether the layer uses the intercept term.
layer_prefix : str (default='dec_')
Prefix of each layer name.
Returns
-------
output : tf.Tensor
Reconstructed output of size `output_shape`
"""
x = tf.reshape(x, shape=[-1, np.prod(x.shape[1:])])
x = _dense(x,
hidden_units,
activation,
use_bias=use_bias,
name=f'{layer_prefix}fc_1')
x = _dense(x,
hidden_units,
activation,
use_bias=use_bias,
name=f'{layer_prefix}fc_2')
x = _dense(x,
np.prod(output_shape),
activation=None,
use_bias=use_bias,
name=f'{layer_prefix}output')
output = tf.reshape(x, shape=[-1] + list(output_shape))
return output
def fc_encoder(input_tensor,
hidden_units=1200,
activation='relu',
use_bias=False,
layer_prefix='enc_'):
"""MLP encoder as proposed in section A.4 of "beta-VAE: Learning Basic
Visual Concepts with a Constrained Variational Framework"
(https://openreview.net/references/pdf?id=Sy2fzU9gl).
Parameters
----------
input_tensor : tf.Tensor
Input tensor to encode.
hidden_units : int (default=1200)
Number of hidden units in each FC layer.
activation : str, tf.keras.layers.Activation, tf.nn (default='relu')
Activation function applied to linear dense layer.
use_bias : bool (default=False)
Boolean whether the layer uses the intercept term.
layer_prefix : str (default='enc_')
Prefix of each layer name.
Returns
-------
x : tf.Tensor
Output tensor of the fully connected encoder.
"""
x = tf.reshape(latent_tensor,
shape=[-1, np.prod(input_tensor.shape[1:])],
name='flatten')
x = _dense(x,
hidden_units,
activation,
use_bias=use_bias,
name=f'{layer_prefix}fc_1')
x = _dense(x,
hidden_units,
activation,
use_bias=use_bias,
name=f'{layer_prefix}fc_2')
return x
def conv_decoder(latent_tensor, output_shape, layer_prefix='dec_'):
"""Convolutional network for a Bernoulli decoder as proposed in section A.1
of "beta-VAE: Learning Basic Visual Concepts with a Constrained Variational
Framework" (https://openreview.net/references/pdf?id=Sy2fzU9gl).
Parameters
----------
input_tensor : tf.Tensor
Input tensor to encode.
output_shape : tuple
Shape of input to reconstruct.
layer_prefix : str (default='dec_')
Prefix of each layer name.
Returns
-------
output : tf.Tensor
Output tensor of the conv decoder of shape `output_shape`.
"""
x = tf.reshape(latent_tensor,
shape=[-1, np.prod(latent_tensor.shape[1:])],
name=f'{layer_prefix}flatten_in')
x = _dense(x,
256,
activation='relu',
use_bias=False,
name=f'{layer_prefix}fc_1')
x = _deconv2d(input_tensor,
filters=64,
kernel_size=4,
strides=2,
name=f'{layer_prefix}deconv2d_1')
x = _deconv2d(x,
filters=64,
kernel_size=4,
strides=2,
name=f'{layer_prefix}conv2d_2')
x = _deconv2d(x,
filters=32,
kernel_size=4,
strides=2,
name=f'{layer_prefix}conv2d_3')
x = _deconv2d(x,
filters=32,
kernel_size=4,
strides=2,
name=f'{layer_prefix}conv2d_4')
output = tf.reshape(x, shape=[-1] + list(output_shape))
return x
def conv_encoder(input_tensor, layer_prefix='enc_'):
"""Convolutional network for a Gaussian encoder as proposed in section A.1
of "beta-VAE: Learning Basic Visual Concepts with a Constrained Variational
Framework" (https://openreview.net/references/pdf?id=Sy2fzU9gl).
Parameters
----------
input_tensor : tf.Tensor
Input tensor to encode.
layer_prefix : str (default='enc_')
Prefix of each layer name.
Returns
-------
x : tf.Tensor
Output tensor of the encoder.
"""
x = _conv2d_bn(input_tensor,
filters=32,
kernel_size=4,
strides=2,
name=f'{layer_prefix}conv2d_1')
x = _conv2d_bn(x,
filters=32,
kernel_size=4,
strides=2,
name=f'{layer_prefix}conv2d_2')
x = _conv2d_bn(x,
filters=64,
kernel_size=4,
strides=2,
name=f'{layer_prefix}conv2d_3')
x = _conv2d_bn(x,
filters=64,
kernel_size=4,
strides=2,
name=f'{layer_prefix}conv2d_4')
x = tf.keras.layers.Flatten(name=f'{layer_prefix}flatten')(x)
x = _dense(x,
out_units=256,
activation='relu',
use_bias=False,
name=f'{layer_prefix}fc')
return x
def factor_discriminator(latent_tensor,
hidden_units=1000,
activation='leaky_relu',
use_bias=False,
layer_prefix='disc_'):
"""Fully connected discriminator as proposed in "Disentangling
by Factorising" (https://arxiv.org/pdf/1802.05983.pdf).
The objective of the discriminator is to determine if the input is sampled
from the true distribtuion q(z) or from a randomly shuffled distribution
q~(z).
Parameters
----------
latent_tensor : tf.Tensor
Tensor from the sampled latent space q(z|x).
hidden_units : int (default=1000)
Number of hidden units in each FC layer.
activation : str, tf.keras.layers.Activation, tf.nn (default='leaky_relu')
Activation function applied to output.
use_bias : bool (default=False)
Boolean whether the layer uses the intercept term.
layer_prefix : str (default='disc_')
Prefix of each layer name.
Returns
-------
logits : tf.Tensor
Logit outputs from the discriminator of size (batch_size, 2).
probs : tf.Tensor
Probability outputs from the discriminiator of size (bathc_size, 2).
"""
if hasattr(tf.nn, activation):
activation = getattr(tf.nn, activation)
x = tf.reshape(latent_tensor, shape=[-1, np.prod(latent_tensor.shape[1:])])
x = _dense(x,
hidden_units,
activation,
use_bias=use_bias,
name=f'{layer_prefix}fc_1')
x = _dense(x,
hidden_units,
activation,
use_bias=use_bias,
name=f'{layer_prefix}fc_2')
x = _dense(x,
hidden_units,
activation,
use_bias=use_bias,
name=f'{layer_prefix}fc_3')
x = _dense(x,
hidden_units,
activation,
use_bias=use_bias,
name=f'{layer_prefix}fc_4')
x = _dense(x,
hidden_units,
activation,
use_bias=use_bias,
name=f'{layer_prefix}fc_5')
x = _dense(x,
hidden_units,
activation,
use_bias=use_bias,
name=f'{layer_prefix}fc_6')
logits = _dense(x,
out_units=2,
activation=None,
use_bias=False,
name=f'{layer_prefix}logit')
probs = tf.nn.softmax(logits, name=f'{layer_prefix}probs')
# without clipping the discriminator diverges (5 hours of debuging)
probs = tf.clip_by_value(probs, 1e-6, 1 - 1e-6)
return logits, probs