def calculate_loss(self, x, beta=1., average=False):
# pass through VAE
x_mean, x_logvar, z1_q, z1_q_mean, z1_q_logvar, z2_q, z2_q_mean, z2_q_logvar, z1_p_mean, z1_p_logvar = self.forward(
x)
# RE
if self.args.input_type == 'binary':
RE = log_Bernoulli(x, x_mean, dim=1)
elif self.args.input_type == 'gray' or self.args.input_type == 'continuous':
RE = -log_Logistic_256(x, x_mean, x_logvar, dim=1)
else:
raise Exception('Wrong input type!')
# KL
log_p_z1 = log_Normal_diag(z1_q, z1_p_mean, z1_p_logvar, dim=1)
log_q_z1 = log_Normal_diag(z1_q, z1_q_mean, z1_q_logvar, dim=1)
log_p_z2 = self.log_p_z2(z2_q)
log_q_z2 = log_Normal_diag(z2_q, z2_q_mean, z2_q_logvar, dim=1)
KL = -(log_p_z1 + log_p_z2 - log_q_z1 - log_q_z2)
# full loss
loss = -RE + beta * KL
if average:
loss = torch.mean(loss)
RE = torch.mean(RE)
KL = torch.mean(KL)
return loss, RE, KL
def calculate_loss(self, x, beta=1., average=False):
'''
:param x: input image(s)
:param beta: a hyperparam for warmup
:param average: whether to average loss or not
:return: value of a loss function
'''
# pass through VAE
x_mean, x_logvar, z1_q, z1_q_mean, z1_q_logvar, z2_q, z2_q_mean, z2_q_logvar, z1_p_mean, z1_p_logvar = self.forward(
x)
# RE
if self.args.input_type == 'binary':
RE = log_Bernoulli(x, x_mean, dim=1)
elif self.args.input_type == 'gray' or self.args.input_type == 'continuous':
RE = -log_Logistic_256(x, x_mean, x_logvar, dim=1)
else:
raise Exception('Wrong input type!')
# KL
log_p_z1 = log_Normal_diag(z1_q, z1_p_mean, z1_p_logvar, dim=1)
log_q_z1 = log_Normal_diag(z1_q, z1_q_mean, z1_q_logvar, dim=1)
log_p_z2 = self.log_p_z2(z2_q)
log_q_z2 = log_Normal_diag(z2_q, z2_q_mean, z2_q_logvar, dim=1)
KL = -(log_p_z1 + log_p_z2 - log_q_z1 - log_q_z2)
loss = -RE + beta * KL
if average:
loss = torch.mean(loss)
RE = torch.mean(RE)
KL = torch.mean(KL)
return loss, RE, KL
def log_p_z2(self, z2):
if self.args.prior == 'standard':
log_prior = log_Normal_standard(z2, dim=1)
elif self.args.prior == 'vampprior':
# z - MB x M
C = self.args.number_components
# calculate params
X = self.means(self.idle_input).view(-1, *self.args.input_size)
# calculate params for given data
z2_p_mean, z2_p_logvar = self.q_z2(X) # C x M)
# expand z
z_expand = z2.unsqueeze(1)
means = z2_p_mean.unsqueeze(0)
logvars = z2_p_logvar.unsqueeze(0)
a = log_Normal_diag(z_expand, means, logvars, dim=2) - math.log(
C) # MB x C
a_max, _ = torch.max(a, 1) # MB
# calculte log-sum-exp
log_prior = (
a_max +
torch.log(torch.sum(torch.exp(a - a_max.unsqueeze(1)), 1))
) # MB
else:
raise Exception('Wrong name of the prior!')
return log_prior