def _log_prob(self, value: tf.Tensor):
""" expect `value` is output from ELU function """
params = self._params
transformed_value, value = _switch_domain(
value,
inputs_domain=self.inputs_domain,
low=self.low,
high=self.high)
## prepare the parameters
if self.n_channels == 1:
component_logits, locs, scales = params
else:
channel_tensors = tf.split(transformed_value,
self.n_channels,
axis=-1)
# If there is more than one channel, we create a linear autoregressive
# dependency among the location parameters of the channels of a single
# pixel (the scale parameters within a pixel are independent). For a pixel
# with R/G/B channels, the `r`, `g`, and `b` saturation values are
# distributed as:
#
# r ~ Logistic(loc_r, scale_r)
# g ~ Logistic(coef_rg * r + loc_g, scale_g)
# b ~ Logistic(coef_rb * r + coef_gb * g + loc_b, scale_b)
component_logits, locs, scales, coeffs = params
loc_tensors = tf.split(locs, self.n_channels, axis=-1)
coef_tensors = tf.split(coeffs, self.n_coeffs, axis=-1)
coef_count = 0
for i in range(self.n_channels):
channel_tensors[i] = channel_tensors[i][..., tf.newaxis, :]
for j in range(i):
loc_tensors[
i] += channel_tensors[j] * coef_tensors[coef_count]
coef_count += 1
locs = tf.concat(loc_tensors, axis=-1)
## create the distrubtion
mixture_distribution = Categorical(logits=component_logits)
# Convert distribution parameters for pixel values in
# `[self._low, self._high]` for use with `QuantizedDistribution`
locs = self.low + 0.5 * (self.high - self.low) * (locs + 1.)
scales = scales * 0.5 * (self.high - self.low)
logistic_dist = QuantizedDistribution(
distribution=TransformedDistribution(
distribution=Logistic(loc=locs, scale=scales),
bijector=Shift(shift=tf.cast(-0.5, self.dtype))),
low=self.low,
high=self.high,
)
dist = MixtureSameFamily(mixture_distribution=mixture_distribution,
components_distribution=Independent(
logistic_dist,
reinterpreted_batch_ndims=1))
dist = Independent(dist, reinterpreted_batch_ndims=2)
return dist.log_prob(value)
def _mean(self, **kwargs):
params = self._params
if self.n_channels == 1:
component_logits, locs, scales = params
else:
# r ~ Logistic(loc_r, scale_r)
# g ~ Logistic(coef_rg * r + loc_g, scale_g)
# b ~ Logistic(coef_rb * r + coef_gb * g + loc_b, scale_b)
component_logits, locs, scales, coeffs = params
loc_tensors = tf.split(locs, self.n_channels, axis=-1)
coef_tensors = tf.split(coeffs, self.n_coeffs, axis=-1)
coef_count = 0
for i in range(self.n_channels):
for j in range(i):
loc_tensors[i] += loc_tensors[j] * coef_tensors[coef_count]
coef_count += 1
locs = tf.concat(loc_tensors, axis=-1)
## create the distrubtion
mixture_distribution = Categorical(logits=component_logits)
# Convert distribution parameters for pixel values in
# `[self._low, self._high]` for use with `QuantizedDistribution`
locs = self.low + 0.5 * (self.high - self.low) * (locs + 1.)
scales = scales * 0.5 * (self.high - self.low)
logistic_dist = TransformedDistribution(
distribution=Logistic(loc=locs, scale=scales),
bijector=Shift(shift=tf.cast(-0.5, self.dtype)))
dist = MixtureSameFamily(mixture_distribution=mixture_distribution,
components_distribution=Independent(
logistic_dist,
reinterpreted_batch_ndims=1))
mean = Independent(dist, reinterpreted_batch_ndims=2).mean()
## normalize the data back to the input domain
return _pixels_to(mean, self.inputs_domain, self.low, self.high)
def MixtureQLogistic(
locs: tf.Tensor,
scales: tf.Tensor,
logits: Optional[tf.Tensor] = None,
probs: Optional[tf.Tensor] = None,
batch_ndims: int = 0,
low: int = 0,
bits: int = 8,
name: str = 'MixtureQuantizedLogistic') -> MixtureSameFamily:
""" Mixture of quantized logistic distribution
Parameters
----------
locs : tf.Tensor
locs of all logistics components, shape `[batch_size, n_components, event_size]`
scales : tf.Tensor
scales of all logistics components, shape `[batch_size, n_components, event_size]`
logits, probs : tf.Tensor
probability for the mixture Categorical distribution, shape `[batch_size, n_components]`
low : int, optional
minimum quantized value, by default 0
bits : int, optional
number of bits for quantization, the maximum will be `2^bits - 1`, by default 8
name : str, optional
distribution name, by default 'MixtureQuantizedLogistic'
Returns
-------
MixtureSameFamily
the mixture of quantized logistic distribution
Example
-------
```
d = MixtureQLogistic(np.ones((12, 3, 8)).astype('float32'),
np.ones((12, 3, 8)).astype('float32'),
logits=np.random.rand(12, 3).astype('float32'),
batch_ndims=1)
```
Reference
---------
Salimans, T., Karpathy, A., Chen, X., Kingma, D.P., 2017.
PixelCNN++: Improving the PixelCNN with Discretized Logistic Mixture
Likelihood and Other Modifications. arXiv:1701.05517 [cs, stat].
"""
cats = Categorical(probs=probs, logits=logits)
dists = Logistic(loc=locs, scale=scales)
dists = TransformedDistribution(
distribution=dists, bijector=Shift(shift=tf.cast(-0.5, dists.dtype)))
dists = QuantizedDistribution(dists, low=low, high=2**bits - 1.)
dists = Independent(dists, reinterpreted_batch_ndims=batch_ndims)
dists = MixtureSameFamily(mixture_distribution=cats,
components_distribution=dists,
name=name)
return dists
def test_pad_mixture_dimensions_mixture_same_family(self):
gm = MixtureSameFamily(
mixture_distribution=Categorical(probs=[0.3, 0.7]),
components_distribution=MultivariateNormalDiag(
loc=[[-1., 1], [1, -1]], scale_identity_multiplier=[1.0, 0.5]))
x = tf.constant([[1.0, 2.0], [3.0, 4.0]])
x_pad = distribution_util.pad_mixture_dimensions(
x, gm, gm.mixture_distribution, tensorshape_util.rank(gm.event_shape))
x_out, x_pad_out = self.evaluate([x, x_pad])
self.assertAllEqual(x_pad_out.shape, [2, 2, 1])
self.assertAllEqual(x_out.reshape([-1]), x_pad_out.reshape([-1]))
def set_prior(self,
loc=0.,
log_scale=np.log(np.expm1(1)),
mixture_logits=None):
r""" Set the prior for mixture density network
loc : Scalar or Tensor with shape `[n_components, event_size]`
log_scale : Scalar or Tensor with shape
`[n_components, event_size]` for 'none' and 'diag' component, and
`[n_components, event_size*(event_size +1)//2]` for 'full' component.
mixture_logits : Scalar or Tensor with shape `[n_components]`
"""
event_size = self.event_size
if self.covariance == 'diag':
scale_shape = [self.n_components, event_size]
fn = lambda l, s: MultivariateNormalDiag(
loc=l, scale_diag=tf.nn.softplus(s))
elif self.covariance == 'none':
scale_shape = [self.n_components, event_size]
fn = lambda l, s: Independent(
Normal(loc=l, scale=tf.math.softplus(s)), 1)
elif self.covariance == 'full':
scale_shape = [
self.n_components, event_size * (event_size + 1) // 2
]
fn = lambda l, s: MultivariateNormalTriL(
loc=l,
scale_tril=FillScaleTriL(diag_shift=1e-5)(tf.math.softplus(s)))
#
if isinstance(log_scale, Number) or tf.rank(log_scale) == 0:
loc = tf.fill([self.n_components, self.event_size], loc)
#
if isinstance(log_scale, Number) or tf.rank(log_scale) == 0:
log_scale = tf.fill(scale_shape, log_scale)
#
if mixture_logits is None:
p = 1. / self.n_components
mixture_logits = np.log(p / (1. - p))
if isinstance(mixture_logits, Number) or tf.rank(mixture_logits) == 0:
mixture_logits = tf.fill([self.n_components], mixture_logits)
#
loc = tf.cast(loc, self.dtype)
log_scale = tf.cast(log_scale, self.dtype)
mixture_logits = tf.cast(mixture_logits, self.dtype)
self._prior = MixtureSameFamily(
components_distribution=fn(loc, log_scale),
mixture_distribution=Categorical(logits=mixture_logits),
name="prior")
return self
def model_gmmprior(args: Arguments):
nets = get_networks(args.ds, zdim=args.zdim, is_hierarchical=False,
is_semi_supervised=False)
latent_size = np.prod(nets['latents'].event_shape)
n_components = 100
loc = tf.compat.v1.get_variable(name="loc", shape=[n_components, latent_size])
raw_scale_diag = tf.compat.v1.get_variable(
name="raw_scale_diag", shape=[n_components, latent_size])
mixture_logits = tf.compat.v1.get_variable(
name="mixture_logits", shape=[n_components])
nets['latents'].prior = MixtureSameFamily(
components_distribution=MultivariateNormalDiag(
loc=loc,
scale_diag=tf.nn.softplus(raw_scale_diag) + tf.math.exp(-7.)),
mixture_distribution=Categorical(logits=mixture_logits),
name="prior")
return VariationalAutoencoder(**nets, name='GMMPrior')
def model_fullcovgmm(args: Arguments):
nets = get_networks(args.ds, zdim=args.zdim, is_hierarchical=False,
is_semi_supervised=False)
latent_size = int(np.prod(nets['latents'].event_shape))
n_components = 100
loc = tf.compat.v1.get_variable(name="loc", shape=[n_components, latent_size])
raw_scale_diag = tf.compat.v1.get_variable(
name="raw_scale_diag", shape=[n_components, latent_size])
mixture_logits = tf.compat.v1.get_variable(
name="mixture_logits", shape=[n_components])
nets['latents'] = RVconf(
event_shape=latent_size,
projection=True,
posterior='mvntril',
prior=MixtureSameFamily(
components_distribution=MultivariateNormalDiag(
loc=loc,
scale_diag=tf.nn.softplus(raw_scale_diag) + tf.math.exp(-7.)),
mixture_distribution=Categorical(logits=mixture_logits),
name="prior"),
name='latents').create_posterior()
return VariationalAutoencoder(**nets, name='FullCov')