tfp.distributions.PixelCNN

The Pixel CNN++ distribution.

Inherits From: Distribution

tfp.distributions.PixelCNN(
    image_shape,
    conditional_shape=None,
    num_resnet=5,
    num_hierarchies=3,
    num_filters=160,
    num_logistic_mix=10,
    receptive_field_dims=(3, 3),
    dropout_p=0.5,
    resnet_activation='concat_elu',
    use_weight_norm=True,
    use_data_init=True,
    high=255,
    low=0,
    dtype=tf.float32,
    name='PixelCNN'
)

Pixel CNN++ [(Salimans et al., 2017)][1] models a distribution over image data, parameterized by a neural network. It builds on Pixel CNN and Conditional Pixel CNN, as originally proposed by [(van den Oord et al., 2016)][2, 3]. The model expresses the joint distribution over pixels as the product of conditional distributions: p(x|h) = prod{ p(x[i] | x[0:i], h) : i=0, ..., d }, in which p(x[i] | x[0:i], h) : i=0, ..., d is the probability of the i-th pixel conditional on the pixels that preceded it in raster order (color channels in RGB order, then left to right, then top to bottom). h is optional additional data on which to condition the image distribution, such as class labels or VAE embeddings. The Pixel CNN++ network enforces the dependency structure among pixels by applying a mask to the kernels of the convolutional layers that ensures that the values for each pixel depend only on other pixels up and to the left (see tfd.PixelCnnNetwork).

Pixel values are modeled with a mixture of quantized logistic distributions, which can take on a set of distinct integer values (e.g. between 0 and 255 for an 8-bit image).

Color intensity v of each pixel is modeled as:

v ~ sum{q[i] * quantized_logistic(loc[i], scale[i]) : i = 0, ..., k },

in which k is the number of mixture components and the q[i] are the Categorical probabilities over the components.

Sampling

Pixels are sampled one at a time, in raster order. This enforces the autoregressive dependency structure, in which the sample of pixel i is conditioned on the samples of pixels 1, ..., i-1. A single color image is sampled as follows:

samples = random_uniform([image_height, image_width, image_channels])
for i in image_height:
  for j in image_width:
    component_logits, locs, scales, coeffs = pixel_cnn_network(samples)
    components = Categorical(component_logits).sample()
    locs = gather(locs, components)
    scales = gather(scales, components)

    coef_count = 0
    channel_samples = []
    for k in image_channels:
      loc = locs[k]
      for m in range(k):
        loc += channel_samples[m] * coeffs[coef_count]
        coef_count += 1
      channel_samp = Logistic(loc, scales[k]).sample()
      channel_samples.append(channel_samp)
    samples[i, j, :] = tf.stack(channel_samples, axis=-1)
samples = round(samples)

Examples


# Build a small Pixel CNN++ model to train on MNIST.

import tensorflow as tf
import tensorflow_datasets as tfds
import tensorflow_probability as tfp

tfd = tfp.distributions
tfk = tf.keras
tfkl = tf.keras.layers

# Load MNIST from tensorflow_datasets
data = tfds.load('mnist')
train_data, test_data = data['train'], data['test']

def image_preprocess(x):
  x['image'] = tf.cast(x['image'], tf.float32)
  return (x['image'],)  # (input, output) of the model

batch_size = 16
train_it = train_data.map(image_preprocess).batch(batch_size).shuffle(1000)

image_shape = (28, 28, 1)
# Define a Pixel CNN network
dist = tfd.PixelCNN(
    image_shape=image_shape,
    num_resnet=1,
    num_hierarchies=2,
    num_filters=32,
    num_logistic_mix=5,
    dropout_p=.3,
)

# Define the model input
image_input = tfkl.Input(shape=image_shape)

# Define the log likelihood for the loss fn
log_prob = dist.log_prob(image_input)

# Define the model
model = tfk.Model(inputs=image_input, outputs=log_prob)
model.add_loss(-tf.reduce_mean(log_prob))

# Compile and train the model
model.compile(
    optimizer=tfk.optimizers.Adam(.001),
    metrics=[])

model.fit(train_it, epochs=10, verbose=True)

# sample five images from the trained model
samples = dist.sample(5)

To train a class-conditional model:


data = tfds.load('mnist')
train_data, test_data = data['train'], data['test']

def image_preprocess(x):
  x['image'] = tf.cast(x['image'], tf.float32)
  # return model (inputs, outputs): inputs are (image, label) and there are no
  # outputs
  return ((x['image'], x['label']),)

batch_size = 16
train_ds = train_data.map(image_preprocess).batch(batch_size).shuffle(1000)
optimizer = tfk.optimizers.Adam()

image_shape = (28, 28, 1)
label_shape = ()
dist = tfd.PixelCNN(
    image_shape=image_shape,
    conditional_shape=label_shape,
    num_resnet=1,
    num_hierarchies=2,
    num_filters=32,
    num_logistic_mix=5,
    dropout_p=.3,
)

image_input = tfkl.Input(shape=image_shape)
label_input = tfkl.Input(shape=label_shape)

log_prob = dist.log_prob(image_input, conditional_input=label_input)

class_cond_model = tfk.Model(
    inputs=[image_input, label_input], outputs=log_prob)
class_cond_model.add_loss(-tf.reduce_mean(log_prob))
class_cond_model.compile(
    optimizer=tfk.optimizers.Adam(),
    metrics=[])
class_cond_model.fit(train_ds, epochs=10)

# Take 10 samples of the digit '5'
samples = dist.sample(10, conditional_input=5.)

# Take 4 samples each of the digits '1', '2', '3'.
# Note that when a batch of conditional input is passed, the sample shape
# (the first argument of `dist.sample`) must have its last dimension(s) equal
# the batch shape of the conditional input (here, (3,)).
samples = dist.sample((4, 3), conditional_input=[1., 2., 3.])

References

[1]: Tim Salimans, Andrej Karpathy, Xi Chen, and Diederik P. Kingma. PixelCNN++: Improving the PixelCNN with Discretized Logistic Mixture Likelihood and Other Modifications. In International Conference on Learning Representations, 2017. https://pdfs.semanticscholar.org/9e90/6792f67cbdda7b7777b69284a81044857656.pdf Additional details at https://github.com/openai/pixel-cnn

[2]: Aaron van den Oord, Nal Kalchbrenner, Oriol Vinyals, Lasse Espeholt, Alex Graves, and Koray Kavukcuoglu. Conditional Image Generation with PixelCNN Decoders. In Neural Information Processing Systems, 2016. https://arxiv.org/abs/1606.05328

[3]: Aaron van den Oord, Nal Kalchbrenner, and Koray Kavukcuoglu. Pixel Recurrent Neural Networks. In International Conference on Machine Learning, 2016. https://arxiv.org/pdf/1601.06759.pdf

Args
`image_shape`	3D `TensorShape` or tuple for the `[height, width, channels]` dimensions of the image.
`conditional_shape`	`TensorShape` or tuple for the shape of the conditional input, or `None` if there is no conditional input.
`num_resnet`	`int`, the number of layers (shown in Figure 2 of [2]) within each highest-level block of Figure 2 of [1].
`num_hierarchies`	`int`, the number of hightest-level blocks (separated by expansions/contractions of dimensions in Figure 2 of [1].)
`num_filters`	`int`, the number of convolutional filters.
`num_logistic_mix`	`int`, number of components in the logistic mixture distribution.
`receptive_field_dims`	`tuple`, height and width in pixels of the receptive field of the convolutional layers above and to the left of a given pixel. The width (second element of the tuple) should be odd. Figure 1 (middle) of [2] shows a receptive field of (3, 5) (the row containing the current pixel is included in the height). The default of (3, 3) was used to produce the results in [1].
`dropout_p`	`float`, the dropout probability. Should be between 0 and 1.
`resnet_activation`	`string`, the type of activation to use in the resnet blocks. May be 'concat_elu', 'elu', or 'relu'.
`use_weight_norm`	`bool`, if `True` then use weight normalization (works only in Eager mode).
`use_data_init`	`bool`, if `True` then use data-dependent initialization (has no effect if `use_weight_norm` is `False`).
`high`	`int`, the maximum value of the input data (255 for an 8-bit image).
`low`	`int`, the minimum value of the input data.
`dtype`	Data type of the `Distribution`.
`name`	`string`, the name of the `Distribution`.

Attributes
`allow_nan_stats`	Python `bool` describing behavior when a stat is undefined. Stats return +/- infinity when it makes sense. E.g., the variance of a Cauchy distribution is infinity. However, sometimes the statistic is undefined, e.g., if a distribution's pdf does not achieve a maximum within the support of the distribution, the mode is undefined. If the mean is undefined, then by definition the variance is undefined. E.g. the mean for Student's T for df = 1 is undefined (no clear way to say it is either + or - infinity), so the variance = E[(X - mean)**2] is also undefined.
`batch_shape`	Shape of a single sample from a single event index as a `TensorShape`. May be partially defined or unknown. The batch dimensions are indexes into independent, non-identical parameterizations of this distribution.
`dtype`	The `DType` of `Tensor`s handled by this `Distribution`.
`event_shape`	Shape of a single sample from a single batch as a `TensorShape`. May be partially defined or unknown.
`experimental_shard_axis_names`	The list or structure of lists of active shard axis names.
`name`	Name prepended to all ops created by this `Distribution`.
`name_scope`	Returns a `tf.name_scope` instance for this class.
`non_trainable_variables`	Sequence of non-trainable variables owned by this module and its submodules. Note: this method uses reflection to find variables on the current instance and submodules. For performance reasons you may wish to cache the result of calling this method if you don't expect the return value to change.
`parameters`	Dictionary of parameters used to instantiate this `Distribution`.
`reparameterization_type`	Describes how samples from the distribution are reparameterized. Currently this is one of the static instances `tfd.FULLY_REPARAMETERIZED` or `tfd.NOT_REPARAMETERIZED`.
`submodules`	Sequence of all sub-modules. Submodules are modules which are properties of this module, or found as properties of modules which are properties of this module (and so on). `a = tf.Module()` `b = tf.Module()` `c = tf.Module()` `a.b = b` `b.c = c` `list(a.submodules) == [b, c]` `True` `list(b.submodules) == [c]` `True` `list(c.submodules) == []` `True`
`trainable_variables`	Sequence of trainable variables owned by this module and its submodules. Note: this method uses reflection to find variables on the current instance and submodules. For performance reasons you may wish to cache the result of calling this method if you don't expect the return value to change.
`validate_args`	Python `bool` indicating possibly expensive checks are enabled.
`variables`	Sequence of variables owned by this module and its submodules. Note: this method uses reflection to find variables on the current instance and submodules. For performance reasons you may wish to cache the result of calling this method if you don't expect the return value to change.

Args
`value`	`float` or `double` `Tensor`.
`name`	Python `str` prepended to names of ops created by this function.
`**kwargs`	Named arguments forwarded to subclass implementation.

Args
`other`	`tfp.distributions.Distribution` instance.
`name`	Python `str` prepended to names of ops created by this function.

Args
`*args`	Passed to implementation `_default_event_space_bijector`.
`**kwargs`	Passed to implementation `_default_event_space_bijector`.

Args
`value`	a `Tensor` valid sample from this distribution family.
`sample_ndims`	Positive `int` Tensor number of leftmost dimensions of `value` that index i.i.d. samples. Default value: `1`.
`validate_args`	Python `bool`, default `False`. When `True`, distribution parameters are checked for validity despite possibly degrading runtime performance. When `False`, invalid inputs may silently render incorrect outputs. Default value: `False`.
`**init_kwargs`	Additional keyword arguments passed through to `cls.__init__`. These take precedence in case of collision with the fitted parameters; for example, `tfd.Normal.experimental_fit([1., 1.], scale=20.)` returns a Normal distribution with `scale=20.` rather than the maximum likelihood parameter `scale=0.`.

Args
`value`	`float` or `double` `Tensor`.
`backward_compat`	`bool` specifying whether to fall back to returning `FullSpace` as the tangent space, and representing R^n with the standard basis.
`**kwargs`	Named arguments forwarded to subclass implementation.

Returns
`log_prob`	a `Tensor` representing the log probability density, of shape `sample_shape(x) + self.batch_shape` with values of type `self.dtype`.
`tangent_space`	a `TangentSpace` object (by default `FullSpace`) representing the tangent space to the manifold at `value`.

Args
`sample_shape`	integer `Tensor` desired shape of samples to draw. Default value: `()`.
`seed`	PRNG seed; see `tfp.random.sanitize_seed` for details. Default value: `None`.
`name`	name to give to the op. Default value: `'sample_and_log_prob'`.
`**kwargs`	Named arguments forwarded to subclass implementation.

Returns
`samples`	a `Tensor`, or structure of `Tensor`s, with prepended dimensions `sample_shape`.
`log_prob`	a `Tensor` of shape `sample_shape(x) + self.batch_shape` with values of type `self.dtype`.

Args
`value`	`Tensor` or Numpy array of image data. May have leading batch dimension(s), which must broadcast to the leading batch dimensions of `conditional_input`.
`conditional_input`	`Tensor` on which to condition the distribution (e.g. class labels), or `None`. May have leading batch dimension(s), which must broadcast to the leading batch dimensions of `value`.
`training`	`bool` or `None`. If `bool`, it controls the dropout layer, where `True` implies dropout is active. If `None`, it defaults to `tf.keras.backend.learning_phase()`.

Args
`sample_shape`	`Tensor` or python list/tuple. Desired shape of a call to `sample()`.
`name`	name to prepend ops with.

Args
`dtype`	Optional float `dtype` to assume for continuous-valued parameters. Some constraining bijectors require advance knowledge of the dtype because certain constants (e.g., `tfb.Softplus.low`) must be instantiated with the same dtype as the values to be transformed.
`num_classes`	Optional `int` `Tensor` number of classes to assume when inferring the shape of parameters for categorical-like distributions. Otherwise ignored.

Args
`n`	`int`, number of samples desired.
`seed`	PRNG seed; see `tfp.random.sanitize_seed` for details.
`conditional_input`	`Tensor` on which to condition the distribution (e.g. class labels), or `None`.
`training`	`bool` or `None`. If `bool`, it controls the dropout layer, where `True` implies dropout is active. If `None`, it defers to Keras' handling of train/eval status.

Args
`sample_shape`	0D or 1D `int32` `Tensor`. Shape of the generated samples.
`seed`	PRNG seed; see `tfp.random.sanitize_seed` for details.
`name`	name to give to the op.
`**kwargs`	Named arguments forwarded to subclass implementation.

tfp.distributions.PixelCNN

Sampling

Examples

References

Args

Attributes

Methods

batch_shape_tensor

cdf

copy

covariance

cross_entropy

entropy

event_shape_tensor

experimental_default_event_space_bijector

experimental_fit

experimental_local_measure

experimental_sample_and_log_prob

is_scalar_batch

is_scalar_event

kl_divergence

log_cdf

log_prob

log_survival_function

mean

mode

param_shapes

param_static_shapes

parameter_properties

prob

quantile

sample

stddev

survival_function

unnormalized_log_prob

variance

with_name_scope

__getitem__

__iter__

`batch_shape_tensor`

`cdf`

`copy`

`covariance`

`cross_entropy`

`entropy`

`event_shape_tensor`

`experimental_default_event_space_bijector`

`experimental_fit`

`experimental_local_measure`

`experimental_sample_and_log_prob`

`is_scalar_batch`

`is_scalar_event`

`kl_divergence`

`log_cdf`

`log_prob`

`log_survival_function`

`mean`

`mode`

`param_shapes`

`param_static_shapes`

`parameter_properties`

`prob`

`quantile`

`sample`

`stddev`

`survival_function`

`unnormalized_log_prob`

`variance`

`with_name_scope`

`getitem`

`iter`