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
).
See also
tfsnippet.distributions.Distribution.log_prob()
for- contents about group_ndims.
-
_forward
(inputs, n_z=None, **kwargs)¶ Get a z sample from \(q(z|h(x))\), using the variational net.
Parameters: Returns: The z samples.
Return type:
-
chain
(x, n_z=None)¶ Chain \(q(z|h(x))\) and \(p(x,z|h(x))\) together.
This method chains the variational net \(q(z|h(x))\) and the model net \(p(x,z|h(x))\) together, with specified observation x. It is typically used to derive the training objectives of VAE. It can also be used to calculate the reconstruction probability (“Variational Autoencoder based Anomaly Detection using Reconstruction Probability”, An, J. and Cho, S. 2015) of x.
Notes
The constructed
VariationalChain
have x observed in its model net, thus this method cannot be used to get reconstructed samples. Usereconstruct()
instead to obtain x samples.Parameters: - x – The input observation x.
- n_z – Number of z samples to take. (default
None
)
Returns: The variational chain.
Return type:
-
get_training_loss
(x, n_z=None)¶ Get the training loss for this VAE.
The variational solver is automatically chosen according to z.is_reparameterized, and the argument n_z, by the following rules:
If z.is_reparameterized is
True
, then:- If n_z > 1, use iwae.
- If n_z == 1 or n_z is
None
, use sgvb.
If z.is_reparameterized is
False
, then:- If n_z > 1, use vimco.
- If n_z == 1 or n_z is
None
, use reinforce.
Dynamic n_z is not supported by this method. Also, Reweighted Wake-Sleep algorithm is not a choice of this method. To derive the training loss for either situation, use
chain()
to obtain aVariationalChain
, and further obtain the loss by chain.vi.training.[algorithm].Parameters: Returns: - A 0-d tensor, the training loss which can be optimized
by gradient descent.
Return type: tf.Tensor
-
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: - z – If specified, observe z in the model net. (default
-
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:
-
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:
-
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
andauto_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 scopeself.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 liketf.concat()
ortf.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 ofListMapper
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 scopeself.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 orModule
instances, which produces corresponding outputs given the inputs.A typical usage of
DictMapper
is to derive the distribution parameters for aDistributionFactory
, 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 scopeself.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 withLambda
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 scopeself.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 ofModule
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 thisSequential
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 scopesequential/_0/
, while the parameters of the second hidden layer (derived by an instance oftensorflow.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 wholeSequential
module. The outputs of the last component will be the outputs of theSequential
module. - name (str) – Optional name of this module
(argument of
VarScopeObject
). - scope (str) – Optional scope of this module
(argument of
VarScopeObject
).
- components (list[(inputs, **kwargs) -> outputs]) – Components of
this
-
_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 scopeself.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
orzhusuan.distributions.Distribution
, or an instance ofDistributionFactory
. - name (str) – Name of the argument (in error message).
Returns: - A distribution factory, for constructing the
desired distribution.
Return type: Raises: TypeError
– If neither of the above conditions is satisfied.- factory – A class object which is subclass of