def loss(
self,
mu_qz,
logvar_qz,
mu_qz0,
logvar_qz0,
mu_pz0,
logvar_pz0,
logit_px,
target_x,
beta=1.0,
):
# kld loss: log q(z|z0, x) - log p(z)
kld_loss = loss_kld_gaussian(mu_qz, logvar_qz, do_sum=False)
# aux dec loss: -log r(z0|z,x)
aux_kld_loss = loss_kld_gaussian_vs_gaussian(
mu_qz0,
logvar_qz0,
mu_pz0,
logvar_pz0,
do_sum=False,
)
# recon loss (neg likelihood): -log p(x|z)
recon_loss = loss_recon_bernoulli_with_logit(logit_px,
target_x,
do_sum=False)
# add loss
loss = recon_loss + beta * kld_loss + beta * aux_kld_loss
return loss.mean(), recon_loss.mean(), kld_loss.mean(
), aux_kld_loss.mean()
def loss(self, mu_z, logvar_z, logit_x, target_x, beta=1.0):
# kld loss
kld_loss = loss_kld_gaussian(mu_z, logvar_z, do_sum=False)
# recon loss (likelihood)
recon_loss = loss_recon_bernoulli_with_logit(logit_x,
target_x,
do_sum=False)
# add loss
loss = recon_loss + beta * kld_loss
return loss.mean(), recon_loss.mean(), kld_loss.mean()
def loss(self, mu_z, logvar_z, mu_x, logvar_x, target_x, beta=1.0):
# kld loss
kld_loss = loss_kld_gaussian(mu_z, logvar_z, do_sum=False)
# recon loss (likelihood)
recon_loss = loss_recon_gaussian(mu_x,
logvar_x,
target_x.view(-1, 2),
do_sum=False)
# add loss
loss = recon_loss + beta * kld_loss
return loss.mean(), recon_loss.mean(), kld_loss.mean()