in vae_helpers.py [0:0]
def discretized_mix_logistic_loss(x, l, low_bit=False):
""" log-likelihood for mixture of discretized logistics, assumes the data has been rescaled to [-1,1] interval """
# Adapted from https://github.com/openai/pixel-cnn/blob/master/pixel_cnn_pp/nn.py
xs = [s for s in x.shape] # true image (i.e. labels) to regress to, e.g. (B,32,32,3)
ls = [s for s in l.shape] # predicted distribution, e.g. (B,32,32,100)
nr_mix = int(ls[-1] / 10) # here and below: unpacking the params of the mixture of logistics
logit_probs = l[:, :, :, :nr_mix]
l = torch.reshape(l[:, :, :, nr_mix:], xs + [nr_mix * 3])
means = l[:, :, :, :, :nr_mix]
log_scales = const_max(l[:, :, :, :, nr_mix:2 * nr_mix], -7.)
coeffs = torch.tanh(l[:, :, :, :, 2 * nr_mix:3 * nr_mix])
x = torch.reshape(x, xs + [1]) + torch.zeros(xs + [nr_mix]).to(x.device) # here and below: getting the means and adjusting them based on preceding sub-pixels
m2 = torch.reshape(means[:, :, :, 1, :] + coeffs[:, :, :, 0, :] * x[:, :, :, 0, :], [xs[0], xs[1], xs[2], 1, nr_mix])
m3 = torch.reshape(means[:, :, :, 2, :] + coeffs[:, :, :, 1, :] * x[:, :, :, 0, :] + coeffs[:, :, :, 2, :] * x[:, :, :, 1, :], [xs[0], xs[1], xs[2], 1, nr_mix])
means = torch.cat([torch.reshape(means[:, :, :, 0, :], [xs[0], xs[1], xs[2], 1, nr_mix]), m2, m3], dim=3)
centered_x = x - means
inv_stdv = torch.exp(-log_scales)
if low_bit:
plus_in = inv_stdv * (centered_x + 1. / 31.)
cdf_plus = torch.sigmoid(plus_in)
min_in = inv_stdv * (centered_x - 1. / 31.)
else:
plus_in = inv_stdv * (centered_x + 1. / 255.)
cdf_plus = torch.sigmoid(plus_in)
min_in = inv_stdv * (centered_x - 1. / 255.)
cdf_min = torch.sigmoid(min_in)
log_cdf_plus = plus_in - F.softplus(plus_in) # log probability for edge case of 0 (before scaling)
log_one_minus_cdf_min = -F.softplus(min_in) # log probability for edge case of 255 (before scaling)
cdf_delta = cdf_plus - cdf_min # probability for all other cases
mid_in = inv_stdv * centered_x
log_pdf_mid = mid_in - log_scales - 2. * F.softplus(mid_in) # log probability in the center of the bin, to be used in extreme cases (not actually used in our code)
# now select the right output: left edge case, right edge case, normal case, extremely low prob case (doesn't actually happen for us)
# this is what we are really doing, but using the robust version below for extreme cases in other applications and to avoid NaN issue with tf.select()
# log_probs = tf.select(x < -0.999, log_cdf_plus, tf.select(x > 0.999, log_one_minus_cdf_min, tf.log(cdf_delta)))
# robust version, that still works if probabilities are below 1e-5 (which never happens in our code)
# tensorflow backpropagates through tf.select() by multiplying with zero instead of selecting: this requires use to use some ugly tricks to avoid potential NaNs
# the 1e-12 in tf.maximum(cdf_delta, 1e-12) is never actually used as output, it's purely there to get around the tf.select() gradient issue
# if the probability on a sub-pixel is below 1e-5, we use an approximation based on the assumption that the log-density is constant in the bin of the observed sub-pixel value
if low_bit:
log_probs = torch.where(x < -0.999,
log_cdf_plus,
torch.where(x > 0.999,
log_one_minus_cdf_min,
torch.where(cdf_delta > 1e-5,
torch.log(const_max(cdf_delta, 1e-12)),
log_pdf_mid - np.log(15.5))))
else:
log_probs = torch.where(x < -0.999,
log_cdf_plus,
torch.where(x > 0.999,
log_one_minus_cdf_min,
torch.where(cdf_delta > 1e-5,
torch.log(const_max(cdf_delta, 1e-12)),
log_pdf_mid - np.log(127.5))))
log_probs = log_probs.sum(dim=3) + log_prob_from_logits(logit_probs)
mixture_probs = torch.logsumexp(log_probs, -1)
return -1. * mixture_probs.sum(dim=[1, 2]) / np.prod(xs[1:])