def negative_iwae_bound(self, x, iw):
"""
Computes the Importance Weighted Autoencoder Bound
Additionally, we also compute the ELBO KL and reconstruction terms
Args:
x: tensor: (batch, dim): Observations
iw: int: (): Number of importance weighted samples
Returns:
niwae: tensor: (): Negative IWAE bound
kl: tensor: (): ELBO KL divergence to prior
rec: tensor: (): ELBO Reconstruction term
"""
################################################################################
# TODO: Modify/complete the code here
# Compute niwae (negative IWAE) with iw importance samples, and the KL
# and Rec decomposition of the Evidence Lower Bound
#
# Outputs should all be scalar
################################################################################
batch = x.size()[0]
"""
sample iw z's
for z_i in sample:
find p(z_i, all x)
find q(z_i, x)
average
"""
phi_m, phi_v = self.enc.encode(x) # (batch, z_dim)
phi_m, phi_v = ut.duplicate(phi_m,
iw), ut.duplicate(phi_v,
iw) # (batch*iw, z_dim)
x_iw = ut.duplicate(x, iw)
z_hat = ut.sample_gaussian(phi_m, phi_v) # (batch*iw, z_dim)
log_q_zx = ut.log_normal(z_hat, phi_m, phi_v) # (batch*iw)
log_p_z = ut.log_normal(z_hat, *self.z_prior) # (batch*iw)
log_p_xz = ut.log_bernoulli_with_logits(
x_iw, self.dec.decode(z_hat)) # (batch*iw)
f = lambda x: x.reshape(iw, batch).transpose(1, 0)
log_p_xz, log_q_zx, log_p_z = f(log_p_xz), f(log_q_zx), f(log_p_z)
iwae = ut.log_mean_exp(log_p_xz - log_q_zx + log_p_z, -1)
iwae = iwae.mean(0)
niwae = -iwae
kl = ut.log_mean_exp(log_q_zx - log_p_z, -1)
kl = kl.mean(0)
rec = ut.log_mean_exp(log_p_xz, -1)
rec = -rec.mean(0)
################################################################################
# End of code modification
################################################################################
return niwae, kl, rec
def negative_iwae_bound(self, x, iw):
"""
Computes the Importance Weighted Autoencoder Bound
Additionally, we also compute the ELBO KL and reconstruction terms
Args:
x: tensor: (batch, dim): Observations
iw: int: (): Number of importance weighted samples
Returns:
niwae: tensor: (): Negative IWAE bound
kl: tensor: (): ELBO KL divergence to prior
rec: tensor: (): ELBO Reconstruction term
"""
################################################################################
# TODO: Modify/complete the code here
# Compute niwae (negative IWAE) with iw importance samples, and the KL
# and Rec decomposition of the Evidence Lower Bound
#
# Outputs should all be scalar
################################################################################
q_m, q_v = self.enc.encode(x)
q_m_, q_v_ = ut.duplicate(q_m, rep=iw), ut.duplicate(q_v, rep=iw)
z_given_x = ut.sample_gaussian(q_m_, q_v_)
decoded_bernoulli_logits = self.dec.decode(z_given_x)
#duplicate x
x_dup = ut.duplicate(x, rep=iw)
rec = ut.log_bernoulli_with_logits(x_dup, decoded_bernoulli_logits)
#compute kl
p_m, p_v = torch.zeros(q_m.shape), torch.ones(q_v.shape)
p_m_, p_v_ = ut.duplicate(p_m, iw), ut.duplicate(p_v, iw)
#print("p_m", p_m.shape)
log_q_phi = ut.log_normal(z_given_x, q_m_, q_v_) #encoded distribution
log_p = ut.log_normal(z_given_x, p_m_, p_v_) #prior distribution
kl = log_q_phi - log_p
niwae = rec - kl
#reshape to size (iw, bs) and then sum
niwae = ut.log_mean_exp(niwae.reshape(iw, -1), dim=0)
rec = ut.log_mean_exp(rec, dim=0)
kl = ut.log_mean_exp(kl, dim=0)
niwae = -torch.mean(niwae)
kl = torch.mean(kl)
rec = torch.mean(kl)
################################################################################
# End of code modification
################################################################################
return niwae, kl, rec
def get_nll(self, outputs, targets):
"""
:return: negative log-likelihood of a minibatch
"""
if not self.constant_var:
mean, var = ut.gaussian_parameters_ff(outputs, dim=0)
return -torch.mean(ut.log_normal(targets, mean, var))
else:
var = self.pred_var * torch.ones_like(outputs)
return -torch.mean(ut.log_normal(targets, outputs, var))
def negative_iwae_bound(self, x, iw):
"""
Computes the Importance Weighted Autoencoder Bound
Additionally, we also compute the ELBO KL and reconstruction terms
Args:
x: tensor: (batch, dim): Observations
iw: int: (): Number of importance weighted samples
Returns:
niwae: tensor: (): Negative IWAE bound
kl: tensor: (): ELBO KL divergence to prior
rec: tensor: (): ELBO Reconstruction term
"""
################################################################################
# TODO: Modify/complete the code here
# Compute niwae (negative IWAE) with iw importance samples, and the KL
# and Rec decomposition of the Evidence Lower Bound
#
# Outputs should all be scalar
################################################################################
# Compute the mixture of Gaussian prior
pm, pv = ut.gaussian_parameters(self.z_pre, dim=1)
#
# Generate samples.
qm, qv = self.enc.encode(x)
niwaes = []
recs = []
kls = []
for i in range(iw):
z_sample = ut.sample_gaussian(qm, qv).view(-1, qm.shape[1])
rec = self.dec.decode(z_sample)
logptheta_x_g_z = ut.log_bernoulli_with_logits(x, rec)
logptheta_z = ut.log_normal_mixture(z_sample, pm, pv)
logqphi_z_g_x = ut.log_normal(z_sample, qm, qv)
niwae = logptheta_x_g_z + logptheta_z - logqphi_z_g_x
#
# Normal variables.
rec = -ut.log_bernoulli_with_logits(x, rec)
kl = ut.log_normal(z_sample, qm, qv) - ut.log_normal_mixture(
z_sample, pm, pv)
niwaes.append(niwae)
recs.append(rec)
kls.append(kl)
niwaes = torch.stack(niwaes, -1)
niwae = ut.log_mean_exp(niwaes, -1)
kl = torch.stack(kls, -1)
rec = torch.stack(recs, -1)
################################################################################
# End of code modification
################################################################################
return -niwae.mean(), kl.mean(), rec.mean()
def negative_iwae_bound(self, x, iw):
"""
Computes the Importance Weighted Autoencoder Bound
Additionally, we also compute the ELBO KL and reconstruction terms
Args:
x: tensor: (batch, dim): Observations
iw: int: (): Number of importance weighted samples
Returns:
niwae: tensor: (): Negative IWAE bound
kl: tensor: (): ELBO KL divergence to prior
rec: tensor: (): ELBO Reconstruction term
"""
################################################################################
# TODO: Modify/complete the code here
# Compute niwae (negative IWAE) with iw importance samples, and the KL
# and Rec decomposition of the Evidence Lower Bound
#
# Outputs should all be scalar
################################################################################
X_dupl = ut.duplicate(x, iw) # Input "x" is duplicated "iw" times
(m, v) = self.enc.encode(X_dupl) # compute the encoder outut
z = ut.sample_gaussian(
m, v) # sample a point from the multivariate Gaussian
logits = self.dec.decode(z) # pass the sampled "Z" through the decoder
# Calculate log Prob of the output x_hat given latent z
ln_P_x_z = ut.log_bernoulli_with_logits(X_dupl, logits)
# Calculate log(P(z))
#ln_P_z = -torch.sum(z*z, -1)/2.0
ln_P_z = ut.log_normal(z, self.z_prior_m, self.z_prior_v)
# Calculate log(Q(z | x)), Conditional Prob of Latent given x
#ln_q_z_x = -torch.sum((z-m)*(z-m)/(2.0*v) + torch.log(v), -1)
ln_q_z_x = ut.log_normal(z, m, v)
exponent = ln_P_x_z + ln_P_z - ln_q_z_x
exponent = exponent.reshape(iw, -1)
L_m_x = ut.log_mean_exp(exponent, 0)
niwae = -torch.mean(L_m_x)
kl = torch.tensor(0)
rec = torch.tensor(0)
################################################################################
# End of code modification
################################################################################
return niwae, kl, rec
def negative_iwae_bound(self, x, iw):
"""
Computes the Importance Weighted Autoencoder Bound
Additionally, we also compute the ELBO KL and reconstruction terms
Args:
x: tensor: (batch, dim): Observations
iw: int: (): Number of importance weighted samples
Returns:
niwae: tensor: (): Negative IWAE bound
kl: tensor: (): ELBO KL divergence to prior
rec: tensor: (): ELBO Reconstruction term
"""
################################################################################
# TODO: Modify/complete the code here
# Compute niwae (negative IWAE) with iw importance samples, and the KL
# and Rec decomposition of the Evidence Lower Bound
#
# Outputs should all be scalar
################################################################################
m, v = self.enc.encode(x)
# m, v -> (batch, dim)
# (batch, dim) -> (batch*iw, dim)
m = ut.duplicate(m, iw)
# (batch, dim) -> (batch*iw, dim)
v = ut.duplicate(v, iw)
# (batch, dim) -> (batch*iw, dim)
x = ut.duplicate(x, iw)
# z -> (batch*iw, dim)
z = ut.sample_gaussian(m, v)
logits = self.dec.decode(z)
kl = ut.log_normal(z, m, v) - ut.log_normal(z, self.z_prior_m,
self.z_prior_v)
rec = -ut.log_bernoulli_with_logits(x, logits)
nelbo = kl + rec
niwae = -ut.log_mean_exp(-nelbo.reshape(iw, -1), dim=0)
niwae, kl, rec = niwae.mean(), kl.mean(), rec.mean()
################################################################################
# End of code modification
################################################################################
return niwae, kl, rec
def get_nll(self, outputs, targets):
"""
:return: negative log-likelihood of a minibatch
"""
if self.likelihood_cost_form == 'mse':
# This method is not validated
return self.mse_fn(outputs, targets)
elif self.likelihood_cost_form == 'gaussian':
if not self.constant_var:
mean, var = ut.gaussian_parameters(outputs, dim=-1)
return -torch.mean(ut.log_normal(targets, mean, var))
else:
var = self.pred_var * torch.ones_like(outputs)
return -torch.mean(ut.log_normal(targets, outputs, var))
def negative_elbo_bound(self, x, y):
################################################################################
# TODO: Modify/complete the code here
# Compute negative Evidence Lower Bound and its KL and Rec decomposition
#
# Note that we are interested in the ELBO of ln p(x | y)
#
# Note that nelbo = kl + rec
#
# Outputs should all be scalar
################################################################################
m, v = self.enc.encode(x, y)
z = ut.sample_gaussian(m, v)
x_m = self.dec.decode(z, y)
rec = ut.log_normal(x, x_m, self.x_v.expand(x_m.size())).mean()
kl_z = ut.kl_normal(m, v,
self.z_prior_m.expand(m.size()),
self.z_prior_v.expand(v.size())).mean()
nelbo = kl_z - rec
################################################################################
# End of code modification
################################################################################
return nelbo, kl_z, rec
def negative_elbo_bound(self, x, y):
################################################################################
# TODO: Modify/complete the code here
# Compute negative Evidence Lower Bound and its KL and Rec decomposition
#
# Note that we are interested in the ELBO of ln p(x | y)
#
# Note that nelbo = kl + rec
#
# Outputs should all be scalar
################################################################################
# sample z
m, v = self.enc.encode(x, y)
z = ut.sample_gaussian(m, v)
# generate x given z,y
x_logits = self.dec.decode(z, y)
# kl on q(z)
kl_z = ut.kl_normal(m, v, self.z_prior[0], self.z_prior[1])
rec_loss = -ut.log_normal(x, x_logits, 0.1 * torch.ones_like(x_logits))
kl_z, rec_loss = rec_loss.mean(), kl_z.mean()
nelbo = rec_loss + kl_z
################################################################################
# End of code modification
################################################################################
return nelbo, kl_z, rec_loss
def negative_elbo_bound(self, x):
"""
Computes the Evidence Lower Bound, KL and, Reconstruction costs
Args:
x: tensor: (batch, dim): Observations
Returns:
nelbo: tensor: (): Negative evidence lower bound
kl: tensor: (): ELBO KL divergence to prior
rec: tensor: (): ELBO Reconstruction term
"""
print(x)
m, v = self.enc.encode(x)
z = ut.sample_gaussian(m, v)
logits = self.dec.decode(z)
batch_size, dim = m.shape
# Compute KL term
# km = torch.zeros(batch_size,self.k,dim)
# kv = torch.ones(batch_size,self.k,dim)
km = self.km.repeat(batch_size, 1, 1)
kv = self.kv.repeat(batch_size, 1, 1)
kl_vec = ut.log_normal(z, m, v) - ut.log_normal_mixture(z, km, kv)
kl = torch.mean(kl_vec)
# Compute reconstruction loss
rec_vec = torch.neg(ut.log_bernoulli_with_logits(x, logits))
rec = torch.mean(rec_vec)
# Compute nelbo
nelbo = rec + kl
return nelbo, kl, rec
def negative_elbo_bound(self, x, y):
################################################################################
# TODO: Modify/complete the code here
# Compute negative Evidence Lower Bound and its KL and Rec decomposition
#
# Note that we are interested in the ELBO of ln p(x | y)
#
# Note that nelbo = kl + rec
#
# Outputs should all be scalar
################################################################################
q_mu, q_var = self.enc.encode(x, y)
z_samp = ut.sample_gaussian(q_mu, q_var)
logits = self.dec.decode(z_samp, y)
rec = -torch.mean(
ut.log_normal(x, logits, 0.1 * torch.ones_like(logits)))
kl_z = torch.mean(
ut.kl_normal(q_mu, q_var, torch.zeros_like(q_mu),
torch.ones_like(q_var)))
nelbo = kl_z + rec
################################################################################
# End of code modification
################################################################################
return nelbo, kl_z, rec
def kl_elem(self, z, qm, qv):
# Compute the mixture of Gaussian prior
prior_m, prior_v = ut.gaussian_parameters(self.z_pre, dim=1)
log_prob_net = ut.log_normal(z, qm, qv)
log_prob_prior = ut.log_normal_mixture(z, prior_m, prior_v)
# print("log_prob_net:", log_prob_net.mean(), "log_prob_prior:", log_prob_prior.mean())
kl_elem = log_prob_net - log_prob_prior
return kl_elem
def negative_iwae_bound(self, x, iw):
"""
Computes the Importance Weighted Autoencoder Bound
Additionally, we also compute the ELBO KL and reconstruction terms
Args:
x: tensor: (batch, dim): Observations
iw: int: (): Number of importance weighted samples
Returns:
niwae: tensor: (): Negative IWAE bound
kl: tensor: (): ELBO KL divergence to prior
rec: tensor: (): ELBO Reconstruction term
"""
################################################################################
# TODO: Modify/complete the code here
# Compute niwae (negative IWAE) with iw importance samples, and the KL
# and Rec decomposition of the Evidence Lower Bound
#
# Outputs should all be scalar
################################################################################
niwae = 0
for i in range(x.size()[0]):
x_i = x[i][:].view(1, x.size()[1])
x_i = ut.duplicate(x_i, iw)
m, v = self.enc.encode(x_i)
z = ut.sample_gaussian(m, v)
x_hat = self.dec.decode(z)
exponent = ut.log_bernoulli_with_logits(x_i, x_hat) + \
ut.log_normal(z, self.z_prior_m.expand(m.size()), self.z_prior_v.expand(v.size())) \
- ut.log_normal(z, m, v)
niwae += -ut.log_mean_exp(exponent, 0).squeeze()
#print(np.std(exponent.data.cpu().numpy()))
#print(exponent.data.cpu().numpy().shape)
niwae = niwae / x.size()[0]
kl = rec = torch.tensor(0)
################################################################################
# End of code modification
################################################################################
return niwae, kl, rec
def negative_iwae_bound(self, x, iw):
"""
Computes the Importance Weighted Autoencoder Bound
Additionally, we also compute the ELBO KL and reconstruction terms
Args:
x: tensor: (batch, dim): Observations
iw: int: (): Number of importance weighted samples
Returns:
niwae: tensor: (): Negative IWAE bound
kl: tensor: (): ELBO KL divergence to prior
rec: tensor: (): ELBO Reconstruction term
"""
################################################################################
# TODO: Modify/complete the code here
# Compute niwae (negative IWAE) with iw importance samples, and the KL
# and Rec decomposition of the Evidence Lower Bound
#
# Outputs should all be scalar
################################################################################
# Compute the mixture of Gaussian prior
prior = ut.gaussian_parameters(self.z_pre, dim=1)
m, v = self.enc.encode(x)
dist = Normal(loc=m, scale=torch.sqrt(v))
z_sample = dist.rsample(sample_shape=torch.Size([iw]))
log_batch_z_sample_prob = []
kl_batch_z_sample = []
for i in range(iw):
recon_logits = self.dec.decode(z_sample[i])
log_batch_z_sample_prob.append(
ut.log_bernoulli_with_logits(
x, recon_logits)) # [batch, z_sample]
kl_batch_z_sample.append(
ut.log_normal(z_sample[i], m, v) -
ut.log_normal_mixture(z_sample[i], prior[0], prior[1]))
log_batch_z_sample_prob = torch.stack(log_batch_z_sample_prob, dim=1)
kl_batch_z_sample = torch.stack(kl_batch_z_sample, dim=1)
niwae = -ut.log_mean_exp(log_batch_z_sample_prob - kl_batch_z_sample,
dim=1).mean(dim=0)
rec = -torch.mean(log_batch_z_sample_prob, dim=0)
kl = torch.mean(kl_batch_z_sample, dim=0)
################################################################################
# End of code modification
################################################################################
return niwae, kl, rec
def negative_iwae_bound(self, x, iw):
"""
Computes the Importance Weighted Autoencoder Bound
Additionally, we also compute the ELBO KL and reconstruction terms
Args:
x: tensor: (batch, dim): Observations
iw: int: (): Number of importance weighted samples
Returns:
niwae: tensor: (): Negative IWAE bound
kl: tensor: (): ELBO KL divergence to prior
rec: tensor: (): ELBO Reconstruction term
"""
m, v = self.enc.encode(x)
batch_size, dim = m.shape
# Duplicate
m = ut.duplicate(m, iw)
v = ut.duplicate(v, iw)
x = ut.duplicate(x, iw)
z = ut.sample_gaussian(m, v)
logits = self.dec.decode(z)
km = self.km.repeat(batch_size, 1, 1)
kv = self.kv.repeat(batch_size, 1, 1)
km = ut.duplicate(km, iw)
kv = ut.duplicate(kv, iw)
kl_vec = ut.log_normal(z, m, v) - ut.log_normal_mixture(z, km, kv)
kl = torch.mean(kl_vec)
# TODO: compute the values below
rec_vec = ut.log_bernoulli_with_logits(x, logits)
rec = torch.neg(torch.mean(rec_vec))
if iw > 1:
iwtensor = torch.zeros(iw)
j = 0
while j < iw:
i = 0
sum = 0
while i < batch_size:
sum += rec_vec[j * batch_size + i]
i += 1
iwtensor[j] = sum / batch_size - kl
j += 1
niwae = torch.neg(ut.log_mean_exp(iwtensor, 0))
else:
niwae = rec + kl
return niwae, kl, rec
def negative_elbo_bound(self, x):
"""
Computes the Evidence Lower Bound, KL and, Reconstruction costs
Args:
x: tensor: (batch, dim): Observations
Returns:
nelbo: tensor: (): Negative evidence lower bound
kl: tensor: (): ELBO KL divergence to prior
rec: tensor: (): ELBO Reconstruction term
"""
################################################################################
# TODO: Modify/complete the code here
# Compute negative Evidence Lower Bound and its KL and Rec decomposition
#
# To help you start, we have computed the mixture of Gaussians prior
# prior = (m_mixture, v_mixture) for you, where
# m_mixture and v_mixture each have shape (1, self.k, self.z_dim)
#
# Note that nelbo = kl + rec
#
# Outputs should all be scalar
################################################################################
# Compute the mixture of Gaussian prior
(m, v) = self.enc.encode(x) # compute the encoder output
#print(" ***** \n")
#print("x xhape ", x.shape)
#print("m and v shapes = ", m.shape, v.shape)
prior = ut.gaussian_parameters(self.z_pre, dim=1)
#print("prior shapes = ", prior[0].shape, prior[1].shape)
z = ut.sample_gaussian(m, v) # sample a point from the multivariate Gaussian
#print("shape of z = ",z.shape)
logits = self.dec.decode(z) # pass the sampled "Z" through the decoder
#print("logits shape = ", logits.shape)
rec = -torch.mean(ut.log_bernoulli_with_logits(x, logits), -1) # Calculate log Prob of the output
log_prob = ut.log_normal(z, m, v)
log_prob -= ut.log_normal_mixture(z, prior[0], prior[1])
kl = torch.mean(log_prob)
rec = torch.mean(rec)
nelbo = kl + rec
################################################################################
# End of code modification
################################################################################
return nelbo, kl, rec
def negative_elbo_bound(self, x):
"""
Computes the Evidence Lower Bound, KL and, Reconstruction costs
Args:
x: tensor: (batch, dim): Observations
Returns:
nelbo: tensor: (): Negative evidence lower bound
kl: tensor: (): ELBO KL divergence to prior
rec: tensor: (): ELBO Reconstruction term
"""
################################################################################
# TODO: Modify/complete the code here
# Compute negative Evidence Lower Bound and its KL and Rec decomposition
#
# To help you start, we have computed the mixture of Gaussians prior
# prior = (m_mixture, v_mixture) for you, where
# m_mixture and v_mixture each have shape (1, self.k, self.z_dim)
#
# Note that nelbo = kl + rec
#
# Outputs should all be scalar
################################################################################
# Compute the mixture of Gaussian prior
prior = ut.gaussian_parameters(self.z_pre, dim=1)
q_m, q_v = self.enc.encode(x)
#print("q_m", q_m.size())
z_given_x = ut.sample_gaussian(q_m, q_v)
decoded_bernoulli_logits = self.dec.decode(z_given_x)
rec = -ut.log_bernoulli_with_logits(x, decoded_bernoulli_logits)
#rec = -torch.mean(rec)
#terms for KL divergence
log_q_phi = ut.log_normal(z_given_x, q_m, q_v)
#print("log_q_phi", log_q_phi.size())
log_p_theta = ut.log_normal_mixture(z_given_x, prior[0], prior[1])
#print("log_p_theta", log_p_theta.size())
kl = log_q_phi - log_p_theta
#print("kl", kl.size())
nelbo = torch.mean(kl + rec)
rec = torch.mean(rec)
kl = torch.mean(kl)
################################################################################
# End of code modification
################################################################################
return nelbo, kl, rec
def negative_elbo_bound(self, x):
"""
Computes the Evidence Lower Bound, KL and, Reconstruction costs
Args:
x: tensor: (batch, dim): Observations
Returns:
nelbo: tensor: (): Negative evidence lower bound
kl: tensor: (): ELBO KL divergence to prior
rec: tensor: (): ELBO Reconstruction term
"""
################################################################################
# TODO: Modify/complete the code here
# Compute negative Evidence Lower Bound and its KL and Rec decomposition
#
# To help you start, we have computed the mixture of Gaussians prior
# prior = (m_mixture, v_mixture) for you, where
# m_mixture and v_mixture each have shape (1, self.k, self.z_dim)
#
# Note that nelbo = kl + rec
#
# Outputs should all be scalar
################################################################################
# Compute the mixture of Gaussian prior
prior = ut.gaussian_parameters(self.z_pre, dim=1)
prior_m, prior_v = prior
batch = x.shape[0]
qm, qv = self.enc.encode(x)
# Now draw Zs from the posterior qm/qv
z = ut.sample_gaussian(qm, qv)
l_posterior = ut.log_normal(z, qm, qv)
multi_m = prior_m.expand(batch, *prior_m.shape[1:])
multi_v = prior_v.expand(batch, *prior_v.shape[1:])
l_prior = ut.log_normal_mixture(z, multi_m, multi_v)
kls = l_posterior - l_prior
kl = torch.mean(kls)
probs = self.dec.decode(z)
recs = ut.log_bernoulli_with_logits(x, probs)
rec = -1.0 * torch.mean(recs)
nelbo = kl + rec
################################################################################
# End of code modification
################################################################################
return nelbo, kl, rec
def negative_iwae_bound(self, x, iw):
"""
Computes the Importance Weighted Autoencoder Bound
Additionally, we also compute the ELBO KL and reconstruction terms
Args:
x: tensor: (batch, dim): Observations
iw: int: (): Number of importance weighted samples
Returns:
niwae: tensor: (): Negative IWAE bound
kl: tensor: (): ELBO KL divergence to prior
rec: tensor: (): ELBO Reconstruction term
"""
################################################################################
# TODO: Modify/complete the code here
# Compute niwae (negative IWAE) with iw importance samples, and the KL
# and Rec decomposition of the Evidence Lower Bound
#
# Outputs should all be scalar
################################################################################
# Compute the mixture of Gaussian prior
prior = ut.gaussian_parameters(self.z_pre, dim=1)
q_m, q_v = self.enc.encode(x)
q_m_, q_v_ = ut.duplicate(q_m, rep=iw), ut.duplicate(q_v, rep=iw)
z_given_x = ut.sample_gaussian(q_m_, q_v_)
decoded_bernoulli_logits = self.dec.decode(z_given_x)
#duplicate x
x_dup = ut.duplicate(x, rep=iw)
rec = ut.log_bernoulli_with_logits(x_dup, decoded_bernoulli_logits)
log_p_theta = ut.log_normal_mixture(z_given_x, prior[0], prior[1])
log_q_phi = ut.log_normal(z_given_x, q_m_, q_v_)
kl = log_q_phi - log_p_theta
niwae = rec - kl
niwae = ut.log_mean_exp(niwae.reshape(iw, -1), dim=0)
niwae = -torch.mean(niwae)
#yay!
################################################################################
# End of code modification
################################################################################
return niwae, kl, rec
def negative_iwae_bound(self, x, iw):
"""
Computes the Importance Weighted Autoencoder Bound
Additionally, we also compute the ELBO KL and reconstruction terms
Args:
x: tensor: (batch, dim): Observations
iw: int: (): Number of importance weighted samples
Returns:
niwae: tensor: (): Negative IWAE bound
kl: tensor: (): ELBO KL divergence to prior
rec: tensor: (): ELBO Reconstruction term
"""
################################################################################
# TODO: Modify/complete the code here
# Compute niwae (negative IWAE) with iw importance samples, and the KL
# and Rec decomposition of the Evidence Lower Bound
#
# Outputs should all be scalar
################################################################################
# Compute the mixture of Gaussian prior
prior = ut.gaussian_parameters(self.z_pre, dim=1)
N_batches, dims = x.size()
x = ut.duplicate(x, iw)
q_mu, q_var = self.enc.encode(x)
z_samp = ut.sample_gaussian(q_mu, q_var)
logits = self.dec.decode(z_samp)
probs = ut.log_bernoulli_with_logits(x, logits)
kl_vals = -ut.log_normal(z_samp, q_mu, q_var) + ut.log_normal_mixture(z_samp, *prior)
probs = probs + kl_vals
niwae = torch.mean(-ut.log_mean_exp(probs.reshape(N_batches, iw), 1))
kl = torch.tensor(0)
rec = torch.tensor(0)
################################################################################
# End of code modification
################################################################################
return niwae, kl, rec
def negative_elbo_bound(self, x):
"""
Computes the Evidence Lower Bound, KL and, Reconstruction costs
Args:
x: tensor: (batch, dim): Observations
Returns:
nelbo: tensor: (): Negative evidence lower bound
kl: tensor: (): ELBO KL divergence to prior
rec: tensor: (): ELBO Reconstruction term
"""
################################################################################
# TODO: Modify/complete the code here
# Compute negative Evidence Lower Bound and its KL and Rec decomposition
#
# To help you start, we have computed the mixture of Gaussians prior
# prior = (m_mixture, v_mixture) for you, where
# m_mixture and v_mixture each have shape (1, self.k, self.z_dim)
#
# Note that nelbo = kl + rec
#
# Outputs should all be scalar
################################################################################
#
# Compute the mixture of Gaussian prior
pm, pv = ut.gaussian_parameters(self.z_pre, dim=1)
#
# Generate samples.
qm, qv = self.enc.encode(x)
z_sample = ut.sample_gaussian(qm, qv)
rec = self.dec.decode(z_sample)
#
# Compute loss.
# KL divergence between the latent distribution and the prior.
rec = -ut.log_bernoulli_with_logits(x, rec)
# kl = ut.kl_normal(qm, qv, pm, pv)
kl = ut.log_normal(z_sample, qm, qv) - ut.log_normal_mixture(
z_sample, pm, pv)
#
# The liklihood of reproducing the sample image given the parameters.
# Would need to take the average of this otherwise.
nelbo = (kl + rec).mean()
# NELBO: 89.24684143066406. KL: 10.346451759338379. Rec: 78.90038299560547
################################################################################
# End of code modification
################################################################################
return nelbo, kl.mean(), rec.mean()
def negative_elbo_bound(self, x):
"""
Computes the Evidence Lower Bound, KL and, Reconstruction costs
Args:
x: tensor: (batch, dim): Observations
Returns:
nelbo: tensor: (): Negative evidence lower bound
kl: tensor: (): ELBO KL divergence to prior
rec: tensor: (): ELBO Reconstruction term
"""
################################################################################
# TODO: Modify/complete the code here
# Compute negative Evidence Lower Bound and its KL and Rec decomposition
#
# To help you start, we have computed the mixture of Gaussians prior
# prior = (m_mixture, v_mixture) for you, where
# m_mixture and v_mixture each have shape (1, self.k, self.z_dim)
#
# Note that nelbo = kl + rec
#
# Outputs should all be scalar
################################################################################
# Compute the mixture of Gaussian prior
prior = ut.gaussian_parameters(self.z_pre, dim=1)
N_samp, dim = x.size()
q_mu, q_var = self.enc.encode(x)
z_samp = ut.sample_gaussian(q_mu, q_var)
logits = self.dec.decode(z_samp)
rec = -torch.mean(ut.log_bernoulli_with_logits(x, logits))
kl = torch.mean(ut.log_normal(z_samp, q_mu, q_var) - ut.log_normal_mixture(z_samp, *prior))
nelbo = kl + rec
################################################################################
# End of code modification
################################################################################
return nelbo, kl, rec
def negative_elbo_bound(self, x):
"""
Computes the Evidence Lower Bound, KL and, Reconstruction costs
Args:
x: tensor: (batch, dim): Observations
Returns:
nelbo: tensor: (): Negative evidence lower bound
kl: tensor: (): ELBO KL divergence to prior
rec: tensor: (): ELBO Reconstruction term
"""
################################################################################
# TODO: Modify/complete the code here
# Compute negative Evidence Lower Bound and its KL and Rec decomposition
#
# Note that nelbo = kl + rec
#
# Outputs should all be scalar
################################################################################
m, v = self.enc.encode(x)
eps = ut.sample_gaussian(self.z_prior_m.expand(m.size()), self.z_prior_v.expand(v.size()))
z_eps = m + eps.mul(v.pow(0.5))
x_hat = self.dec.decode(z_eps)
rec = ut.log_bernoulli_with_logits(x, x_hat).mean()
#kl = ut.kl_normal(m, v, self.z_prior_m.expand(m.size()),
# self.z_prior_v.expand(v.size())).mean()
hq = ut.log_normal(z_eps, m, v).mean()
#hz = ut.log_normal(z_eps, self.z_prior_m.expand(m.size()), self.z_prior_v.expand(v.size())).mean()
kl = hq
nelbo = kl - rec
################################################################################
# End of code modification
################################################################################
return nelbo, kl, rec
def nelbo(self, x, epoch=None):
"""
Computes the Evidence Lower Bound, KL and, Reconstruction costs
Args:
x: tensor: (batch, dim): Observations
Returns:
nelbo: tensor: (): Negative evidence lower bound
kl: tensor: (): ELBO KL divergence to prior
rec: tensor: (): ELBO Reconstruction term
"""
# get dimensions
z_dim = self.z_dim
z_num = self.z_num
b_size = x.size(0)
# sample c and determine parents
c = self.mu.sample()
mask = torch.zeros((z_num, z_num))
mask[torch.tril(torch.ones((z_num, z_num)), -1) == 1] = c
p_num = (mask != 0).sum(dim=0)
# get data encoding
hx = self.gl_enc.encode(x)
# compute log prior and log posterior
# sample z from posterior
logq = torch.zeros(b_size)
logp = torch.zeros(b_size)
z = torch.zeros((b_size, z_dim, z_num))
for n in range(1, z_num + 1):
# if no parents, sample unit gaussian
# else, sample encoded gaussian
if p_num[-n] == 0:
m_n = torch.zeros((b_size, z_dim), requires_grad=False)
v_n = torch.ones((b_size, z_dim), requires_grad=False)
z[:, :, -n] = ut.sample_gaussian(m_n, v_n)
else:
# get parents of z_n
p_n = z.transpose(1, 2).reshape(-1, z_dim * z_num)
# get bottom-up and top-down encoded gaussian parameters
bu_psi_n = self.bu_enc[-n].encode(hx)
td_psi_n = self.td_enc[-n].encode(p_n)
# compute precision-weighted fusion of encoded gaussian parameters
psi_n = self.gaussian_params_fusion(bu_psi_n, td_psi_n)
# sample z_n from posterior
z_n = ut.sample_gaussian(psi_n[0], psi_n[1])
z[:, :, -n] = z_n
# add to log prior and log posterior
logp += ut.log_normal(z_n, td_psi_n[0], td_psi_n[1])
logq += ut.log_normal(z_n, psi_n[0], psi_n[1])
# compute log conditional
logits = self.dec.decode(hx)
# logits = self.dec.decode(z.transpose(1,2).reshape(-1,z_dim*z_num))
logp_cond = ut.log_bernoulli_with_logits(x, logits)
# compute rec and kl terms
rec = -logp_cond.mean()
kl = (logq - logp).mean()
nelbo = rec + kl
# print(nelbo.data, kl.data, rec.data)
# print(c)
return nelbo, kl, rec
def negative_iwae_bound(self, x, iw):
"""
Computes the Importance Weighted Autoencoder Bound
Additionally, we also compute the ELBO KL and reconstruction terms
Args:
x: tensor: (batch, dim): Observations
iw: int: (): Number of importance weighted samples
Returns:
niwae: tensor: (): Negative IWAE bound
kl: tensor: (): ELBO KL divergence to prior
rec: tensor: (): ELBO Reconstruction term
"""
################################################################################
# TODO: Modify/complete the code here
# Compute niwae (negative IWAE) with iw importance samples, and the KL
# and Rec decomposition of the Evidence Lower Bound
#
# Outputs should all be scalar
################################################################################
# Compute the mixture of Gaussian prior
prior = ut.gaussian_parameters(self.z_pre, dim=1)
prior_m, prior_v = prior
batch = x.shape[0]
multi_x = ut.duplicate(x, iw)
qm, qv = self.enc.encode(x)
multi_qm = ut.duplicate(qm, iw)
multi_qv = ut.duplicate(qv, iw)
# z will be (batch*iw x z_dim)
# with sampled z's for a given x non-contiguous!
z = ut.sample_gaussian(multi_qm, multi_qv)
probs = self.dec.decode(z)
recs = ut.log_bernoulli_with_logits(multi_x, probs)
rec = -1.0 * torch.mean(recs)
multi_m = prior_m.expand(batch * iw, *prior_m.shape[1:])
multi_v = prior_v.expand(batch * iw, *prior_v.shape[1:])
z_priors = ut.log_normal_mixture(z, multi_m, multi_v)
x_posteriors = recs
z_posteriors = ut.log_normal(z, multi_qm, multi_qv)
kls = z_posteriors - z_priors
kl = torch.mean(kls)
log_ratios = z_priors + x_posteriors - z_posteriors
# Should be (batch*iw, z_dim), batch ratios non contiguous
unflat_log_ratios = log_ratios.reshape(iw, batch)
niwaes = ut.log_mean_exp(unflat_log_ratios, 0)
niwae = -1.0 * torch.mean(niwaes)
################################################################################
# End of code modification
################################################################################
return niwae, kl, rec
def negative_iwae_bound(self, x, iw):
"""
Computes the Importance Weighted Autoencoder Bound
Additionally, we also compute the ELBO KL and reconstruction terms
Args:
x: tensor: (batch, dim): Observations
iw: int: (): Number of importance weighted samples
Returns:
niwae: tensor: (): Negative IWAE bound
kl: tensor: (): ELBO KL divergence to prior
rec: tensor: (): ELBO Reconstruction term
"""
################################################################################
# TODO: Modify/complete the code here
# Compute niwae (negative IWAE) with iw importance samples, and the KL
# and Rec decomposition of the Evidence Lower Bound
#
# Outputs should all be scalar
################################################################################
## E_{z(1),...z(n)} {log p(x|z) + log p(z) - log q(z|x)}
m, v = self.enc.encode(x)
batch_m = m.unsqueeze(1)
batch_m = batch_m.repeat(1, iw, 1) # dimension (batch, iw, 10)
batch_v = v.unsqueeze(1)
batch_v = batch_v.repeat(1, iw, 1)
batch_x = x.unsqueeze(1)
batch_x = batch_x.repeat(1, iw, 1)
# log p(x|z)
zs = ut.sample_gaussian(batch_m, batch_v)
logits = self.dec.decode(zs)
raw_probs = ut.log_bernoulli_with_logits(batch_x, logits)
pxz = torch.mean(ut.log_mean_exp(raw_probs, dim=-1))
# log p(z)
batch_size = batch_m.shape[0]
batch_z_prior_m = self.z_prior_m.view(1, 1, -1)
batch_z_prior_m = batch_z_prior_m.repeat(batch_size, iw, 1)
batch_z_prior_v = self.z_prior_v.view(1, 1, -1)
batch_z_prior_v = batch_z_prior_v.repeat(batch_size, iw, 1)
pz = ut.log_normal(zs, batch_z_prior_m, batch_z_prior_v) # (batch, iw)
pz = torch.mean(ut.log_mean_exp(pz, dim=-1))
# log q(z|x)
qzx = ut.log_normal(zs, batch_m, batch_v)
qzx = torch.mean(ut.log_mean_exp(qzx, dim=-1))
# print(pxz, pz, qzx)
niwae = -1 * (pxz + pz - qzx)
rec = pxz - pz
kl = pz - qzx
################################################################################
# End of code modification
################################################################################
return niwae, kl, rec
