tfsnippet.modules¶

class tfsnippet.modules.VAE(p_z, p_x_given_z, q_z_given_x, h_for_p_x, h_for_q_z, z_group_ndims=1, x_group_ndims=1, is_reparameterized=None, name=None, scope=None)¶

Bases: tfsnippet.modules.base.Module

A general implementation of variational auto-encoder as module.

The variational auto-encoder (“Auto-Encoding Variational Bayes”, Kingma, D.P. and Welling) is a deep Bayesian network, with observed variable x and latent variable z. The generative process starts from z with prior distribution $p(z)$, following a hidden network $h(z)$, then comes to x with distribution $p(x|h(z))$. To do posterior inference of $p(z|x)$, variational inference techniques are adopted, to train a separated distribution $q(z|h(x))$ ($h(x)$ denoting the hidden network) to approximate $p(z|x)$.

This class provides a general implementation of variational auto-encoder, with customizable $p(z)$, $p(x|h(z))$, $q(z|h(x))$, as well as the hidden networks $h(z)$ and $h(x)$.

For example, to construct a VAE with diagonal Normal z and x:

from tensorflow import keras as K
from tfsnippet.modules import VAE, DictMapper, Sequential
from tfsnippet.distributions import Normal

batch_size = 128
x_dims, z_dims = 100, 10
vae = VAE(
    p_z=Normal(mean=tf.zeros([batch_size, z_dims]),
               std=tf.ones([batch_size, x_dims])),
    p_x_given_z=Normal,
    q_z_given_x=Normal,
    h_for_p_x=Sequential([
        K.layers.Dense(100, activation=tf.nn.relu),
        DictMapper({'mean': K.layers.Dense(x_dims),
                    'logstd': K.layers.Dense(x_dims)})
    ]),
    h_for_q_z=Sequential([
        K.layers.Dense(100, activation=tf.nn.relu),
        DictMapper({'mean': K.layers.Dense(z_dims),
                    'logstd': K.layers.Dense(z_dims)})
    ])
)

To train the vae:

# Automatically derive a single-sample loss.
# Depending on ``z.is_reparameterized``, it might be derived by
# `sgvb` (is_reparameterized == True) or `reinforce` (otherwise).
loss = vae.get_training_loss(x)

# Automatically derive a multi-sample loss.
# Depending on ``z.is_reparameterized``, it might be derived by
# `iwae` (is_reparameterized == True) or `vimco` (otherwise).
loss = vae.get_training_loss(x, n_z=10)

# Or manually derive a reweighted wake-sleep training loss.
# Note the `VariationalTrainingObjectives` produce per-data
# training objectives, instead of a 0-d scalar loss as the
# `VAE.get_training_loss` does.
chain = vae.chain(x, n_z=10)
loss = tf.reduce_mean(chain.vi.training.rws_wake())

To map from x to z:

# use the :class:`Module` interface for one-to-one mapping
z = vae(x)

# use the :class:`Module` interface for multiple `z` samples
z = vae(x, n_z=10)

# or obtain the variational :class:`BayesianNet` with observed `z`
q_net = vae.variational(x, z=observed_z)
z_log_prob = q_net['z'].log_prob()

To reconstruct x:

# use the :meth:`VAE.reconstruct` for obtaining one `x` sample
x_reconstructed = vae.reconstruct(x)

# to obtain multiple `z` samples, and further multiple `x` samples
# (this results in 100 `x` samples for each input `x`)
x_reconstructed = vae.reconstruct(x, n_z=10, n_x=10)

To sample x from prior z or observed z:

# sample multiple prior `z`, then one `x` for each `z`
x = vae.model(n_z=10)['x']

# sample multiple `x` for each observed `z`
x = vae.model(z=observed_z, n_x=10)['x']

__init__(p_z, p_x_given_z, q_z_given_x, h_for_p_x, h_for_q_z, z_group_ndims=1, x_group_ndims=1, is_reparameterized=None, name=None, scope=None)¶

Construct the VAE.

Parameters:

p_z (Distribution) – $p(z)$, the distribution instance.
p_x_given_z – $p(x|h(z))$, a distribution class or a DistributionFactory object.
q_z_given_x – $q(z|h(x))$, a distribution class or a DistributionFactory object.
h_for_p_x (Module) – $h(z)$, the hidden network module for $p(x|h(z))$. The output of h_for_p_x must be a dict[str, any], the parameters for p_x_given_z.
h_for_q_z (Module) – $h(x)$, the hidden network module for $q(z|h(x))$. The output of h_for_q_z must be a dict[str, any], the parameters for q_z_given_x.
z_group_ndims (int or tf.Tensor) – group_ndims for z. (default 1)
x_group_ndims (int or tf.Tensor) – group_ndims for x. (default 1)
is_reparameterized (bool or None) – Whether or not z should be re-parameterized? (default None, following the settings of z distributions.)
name (str) – Optional name of this module (argument of VarScopeObject).
scope (str) – Optional scope of this module (argument of VarScopeObject).

get_training_objective(*args, **kwargs)¶

h_for_p_x¶

Get $h(z)$, the hidden network for $p(x|h(z))$.

Returns:	The hidden network.
Return type:	Module

h_for_q_z¶

Get $h(x)$, the hidden network for $q(z|h(x))$.

Returns:	The hidden network.
Return type:	Module

is_reparameterized¶: Whether or not z is re-parameterized?

model(**kwargs)¶

Derive an instance of $p(x|h(z))$, the model net.

Parameters:	z – If specified, observe z in the model net. (default `None`) x – If specified, observe x in the model net. (default `None`) n_z – The number of z samples to take for each x, if z is not observed. (default `None`, one z sample for each x, without dedicated sampling dimension) It is recommended to specify this argument even if z is observed, to make explicit how many samples are there in the observation. n_x – The number of x samples to take for each z, if x is not observed. (default `None`, one x sample for each z, without dedicated sampling dimension) It is recommended to specify this argument even if x is observed, to make explicit how many samples are there in the observation.
Returns:	The variational net.
Return type:	BayesianNet

p_x_given_z¶

Get the factory for $p(x|h(z))$.

Returns:	The distribution factory.
Return type:	DistributionFactory

p_z¶

Get $p(z)$, the prior distribution of z.

Returns:	The distribution instance.
Return type:	Distribution

q_z_given_x¶

Get the factory for $q(z|h(x))$.

Returns:	The distribution factory.
Return type:	DistributionFactory

reconstruct(x, n_z=None, n_x=None)¶

Sample reconstructed x from $p(x|h(z))$, where z is (are) sampled from $q(z|h(x))$ using the specified observation x.

Parameters:	x – The observation x for $q(z\|h(x))$. n_z – Number of intermediate z samples to take for each input x. n_x – Number of reconstructed x samples to take for each z.
Returns:	The reconstructed samples x.
Return type:	StochasticTensor

variational(**kwargs)¶

Derive an instance of $q(z|h(x))$, the variational net.

Parameters:	x – The observation x for the variational net. z – If specified, observe z in the variational net. (default `None`) n_z – The number of z samples to take for each x, if z is not observed. (default `None`, one sample for each x, without dedicated sampling dimension) It is recommended to specify this argument even if z is observed, to make explicit how many samples are there in the observation.
Returns:	The variational net.
Return type:	BayesianNet

x_group_ndims¶: Get the group_ndims for x.

z_group_ndims¶: Get the group_ndims for z.

class tfsnippet.modules.Module(name=None, scope=None)¶

Bases: tfsnippet.utils.scope.VarScopeObject

Base class for neural network modules.

A neural network module is basically a reusable object that could derive outputs for inputs, using the same, fixed set of parameters every time. This class provides a base for implementing such a reusable module, which adopts VarScopeObject and auto_reuse_variables() to achieve parameter sharing, making it possible to use all other 3rd party libraries with this module interface of TFSnippet.

For example, a reusable dense layer may be implemented as follows:

class Dense(Module):

    def __init__(self, num_outputs, activation_fn, **kwargs):
        super(Dense, self).__init__(**kwargs)
        self.num_outputs = num_outputs
        self.activation_fn = activation_fn

    def _forward(self, inputs, **kwargs):
        return tf.contrib.layers.fully_connected(
            inputs,
            activation_fn=self.activation_fn,
            num_outputs=self.num_outputs
        )

dense = Dense(num_outputs)
y1 = dense(x1)
y2 = dense(x2)  # parameters shared with y1

It is even possible to construct module instances within _forward() of an already established module instance, with parameter sharing working properly, for example:

class ConstantBiasedSoftPlus(Module):

    def __init__(self, num_outputs, biases, **kwargs):
        super(ConstantBiasedSoftPlus, self).__init__(**kwargs)
        self.num_outputs = num_outputs
        self.biases = biases

    def _forward(self, inputs, **kwargs):
        softplus = Dense(self.num_outputs, activation_fn=tf.nn.softplus)
        outputs = softplus(inputs)
        return outputs + self.biases

__call__(...) <==> x(...)¶

_forward(inputs, **kwargs)¶

Derive outputs for specified inputs.

Overriding this in a child class to actually implement the module. This method is not designed for calling directly. Instead, it is designed to be called from __call__(), where the variable scope self.variable_scope + 'forward' is opened, with reusing flag set automatically.

Parameters:	inputs – The input tensor(s). **kwargs – Named arguments for deriving the outputs.
Returns:	The derived outputs.

class tfsnippet.modules.ListMapper(mapper, name=None, scope=None)¶

Bases: tfsnippet.modules.base.Module

Module that maps inputs into a list of outputs.

This module is a branch module, which maps the inputs into a list of outputs, according to the specified mapper list. Each item of mapper should be a function of a Module instance, which produces corresponding output given the inputs.

To merge the outputs of ListMapper, one may use functions like tf.concat() or tf.stack(), e.g.:

net = Sequential([
    ListMapper([...]),
    functools.partial(tf.concat, axis=0),
])

in which the net module first maps the inputs into a list by using ListMapper, then generates the final output by concatenating the outputs of ListMapper along axis-0.

__init__(mapper, name=None, scope=None)¶

Construct the ListMapper.

Parameters:	mapper (list[(inputs, *kwargs) -> outputs]) – The mapper list. name* (str) – Optional name of this module (argument of `VarScopeObject`). scope (str) – Optional scope of this module (argument of `VarScopeObject`).

_forward(inputs, **kwargs)¶

Derive outputs for specified inputs.

Overriding this in a child class to actually implement the module. This method is not designed for calling directly. Instead, it is designed to be called from __call__(), where the variable scope self.variable_scope + 'forward' is opened, with reusing flag set automatically.

Parameters:	inputs – The input tensor(s). **kwargs – Named arguments for deriving the outputs.
Returns:	The derived outputs.

class tfsnippet.modules.DictMapper(mapper, name=None, scope=None)¶

Bases: tfsnippet.modules.base.Module

Module that maps inputs into a dict of outputs.

This module is a branch module, which maps the inputs into a dict of outputs, according to the specified mapper dict. The keys of mapper should be valid Python identifiers, i.e., matching the pattern ^[A-Za-z_][A-Za-z0-9_]*$. The values of mapper should be functions or Module instances, which produces corresponding outputs given the inputs.

A typical usage of DictMapper is to derive the distribution parameters for a DistributionFactory, e.g.:

from tensorflow import keras as K

from tfsnippet.distributions import Normal
from tfsnippet.modules import Sequential, DictMapper

net = Sequential([
    K.layers.Dense(100, activation=tf.nn.relu),
    DictMapper({
        'mean': K.layers.Dense(2),
        'logstd': K.layers.Dense(2),
    })
])
factory = Normal.factory()
distribution = factory(net(x))

In the above example, the net module will produce a dict carrying mean and logstd, consumed by the factory.

__init__(mapper, name=None, scope=None)¶

Construct the DictMapper.

Parameters:	mapper (dict[str, (inputs, *kwargs) -> outputs]) – The mapper dict. name* (str) – Optional name of this module (argument of `VarScopeObject`). scope (str) – Optional scope of this module (argument of `VarScopeObject`).

_forward(inputs, **kwargs)¶

Derive outputs for specified inputs.

Overriding this in a child class to actually implement the module. This method is not designed for calling directly. Instead, it is designed to be called from __call__(), where the variable scope self.variable_scope + 'forward' is opened, with reusing flag set automatically.

Parameters:	inputs – The input tensor(s). **kwargs – Named arguments for deriving the outputs.
Returns:	The derived outputs.

class tfsnippet.modules.Lambda(f, name=None, scope=None)¶

Bases: tfsnippet.modules.base.Module

Wrapping arbitrary function into a neural network Module.

This class wraps an arbitrary function or lambda expression into a neural network Module, reusing the variables created within the specified function.

For example, one may wrap tensorflow.contrib.layers.fully_connected() into a reusable module with Lambda component as follows:

import functools
from tensorflow.contrib import layers

dense = Lambda(
    functools.partial(
        layers.fully_connected,
        num_outputs=100,
        activation_fn=tf.nn.relu
    )
)

__init__(f, name=None, scope=None)¶

Construct the Lambda.

Parameters:	f ((inputs, *kwargs) -> outputs) – The function or lambda expression which derives the outputs. name* (str) – Optional name of this module (argument of `VarScopeObject`). scope (str) – Optional scope of this module (argument of `VarScopeObject`).

_forward(inputs, **kwargs)¶

Derive outputs for specified inputs.

Overriding this in a child class to actually implement the module. This method is not designed for calling directly. Instead, it is designed to be called from __call__(), where the variable scope self.variable_scope + 'forward' is opened, with reusing flag set automatically.

Parameters:	inputs – The input tensor(s). **kwargs – Named arguments for deriving the outputs.
Returns:	The derived outputs.

class tfsnippet.modules.Sequential(components, name=None, scope=None)¶

Bases: tfsnippet.modules.base.Module

Wrapping a sequential of neural network modules as a Module.

This class wraps a sequential of neural network components into a unified Module, reusing the parameters inside each component. Instances of Module and any arbitrary 3rd-party components can be mixed together freely. For example:

from tensorflow import keras as K
from tensorflow.contrib import layers

mlp = Sequential([
    lambda inputs: layers.fully_connected(
        inputs, num_outputs=100, activation_fn=tf.nn.relu),
    K.layers.Dense(100, activation_fn=tf.nn.relu),
    K.layers.Dense(1),
    tf.nn.sigmoid
])

Which builds a multi-layer perceptron, with 2 hidden layers of 100 units and ReLU activation, plus a sigmoid output layer with 1 unit.

Note

If another instance of Module is specified as one component, the variables of that module is managed within its own scope, instead of the scope of this sequential module. On the other hand, if instead a function or a method is provided, it will be called within a child scope of this Sequential module.

As a result, in the above code example, the parameters of the first hidden layer (derived by fully_connected()) would be collected in the scope sequential/_0/, while the parameters of the second hidden layer (derived by an instance of tensorflow.keras.layers.Dense) would be collected in its own variable scope.

__init__(components, name=None, scope=None)¶

Construct the Sequential.

Parameters:

components (list[(inputs, **kwargs) -> outputs]) – Components of this Sequential module, each should be a callable object which consumes the outputs of previous component as inputs. The first component should consume the inputs as well as the named arguments (**kwargs) given to the whole Sequential module. The outputs of the last component will be the outputs of the Sequential module.
name (str) – Optional name of this module (argument of VarScopeObject).
scope (str) – Optional scope of this module (argument of VarScopeObject).

_forward(inputs, **kwargs)¶

Derive outputs for specified inputs.

Overriding this in a child class to actually implement the module. This method is not designed for calling directly. Instead, it is designed to be called from __call__(), where the variable scope self.variable_scope + 'forward' is opened, with reusing flag set automatically.

Parameters:	inputs – The input tensor(s). **kwargs – Named arguments for deriving the outputs.
Returns:	The derived outputs.

tfsnippet.modules.validate_distribution_factory(factory, name)¶

Validate the specified distribution factory argument.

Parameters:	factory – A class object which is subclass of `Distribution` or `zhusuan.distributions.Distribution`, or an instance of `DistributionFactory`. name (str) – Name of the argument (in error message).
Returns:	A distribution factory, for constructing the desired distribution.
Return type:	DistributionFactory
Raises:	`TypeError` – If neither of the above conditions is satisfied.

tfsnippet.modules¶

TFSnippet

Navigation

Related Topics

Parameters:	inputs – The input x. n_z – Number of samples to taken for z. (default `None`) **kwargs – Capturing and ignoring all other parameters. This is the default behavior of a `Module`.
Returns:	The z samples.
Return type:	StochasticTensor

Parameters:	x – The input observation x. n_z – Number of z samples to take. (default `None`)
Returns:	The variational chain.
Return type:	VariationalChain

Parameters:	x – The observation x for \(q(z\|h(x))\). n_z – Number of intermediate z samples to take for each input x. n_x – Number of reconstructed x samples to take for each z.
Returns:	The reconstructed samples x.
Return type:	StochasticTensor