Get the prior KL term for the variational lower-bound, measured in bits-per-dim. This term can't be optimized, as it only depends on the encoder. Returns a batch of (bs,) KL values in bits.
(self, x_start: Tensor)
| 395 | """ VLB """ |
| 396 | |
| 397 | def prior_bpd(self, x_start: Tensor): |
| 398 | """ |
| 399 | Get the prior KL term for the variational lower-bound, measured in |
| 400 | bits-per-dim. This term can't be optimized, as it only depends on the |
| 401 | encoder. Returns a batch of (bs,) KL values in bits. |
| 402 | """ |
| 403 | batch_size = x_start.shape[0] |
| 404 | t = torch.tensor([self.num_timesteps - 1] * batch_size, device=x_start.device) |
| 405 | qt_mean, _, qt_log_variance = self.q_mean_variance(x_start, t) |
| 406 | kl_prior = normal_kl( |
| 407 | mean1=qt_mean, logvar1=qt_log_variance, mean2=0.0, logvar2=0.0 |
| 408 | ) |
| 409 | prior_bpd = kl_prior.mean(dim=[*range(1, len(kl_prior.shape))]) / np.log(2.0) |
| 410 | return prior_bpd |
| 411 | |
| 412 | def _vb_terms_bpd(self, model, x_start, x_t, t, clip_denoised=False, model_kwargs=None): |
| 413 | """ |
no test coverage detected