Get a term for the variational lower-bound. The resulting units are bits (rather than nats, as one might expect). This allows for comparison to other papers. :return: a dict with the following keys: - 'output': a shape [N] tensor of NLLs or KLs.
(self, model, x_start, x_t, t, clip_denoised=False, model_kwargs=None)
| 410 | return prior_bpd |
| 411 | |
| 412 | def _vb_terms_bpd(self, model, x_start, x_t, t, clip_denoised=False, model_kwargs=None): |
| 413 | """ |
| 414 | Get a term for the variational lower-bound. |
| 415 | |
| 416 | The resulting units are bits (rather than nats, as one might expect). |
| 417 | This allows for comparison to other papers. |
| 418 | |
| 419 | :return: a dict with the following keys: |
| 420 | - 'output': a shape [N] tensor of NLLs or KLs. |
| 421 | - 'pred_xstart': the x_0 predictions. |
| 422 | """ |
| 423 | true_mean, _, true_log_variance_clipped = self.q_posterior(x_start=x_start, x_t=x_t, t=t) |
| 424 | |
| 425 | out = self.p_mean_variance( |
| 426 | model, x_t, t, clip_denoised=clip_denoised, model_kwargs=model_kwargs |
| 427 | ) |
| 428 | kl = normal_kl( |
| 429 | true_mean, true_log_variance_clipped, out["mean"], out["log_variance"] |
| 430 | ) |
| 431 | kl = mean_flat(kl) / np.log(2.0) |
| 432 | |
| 433 | decoder_nll = -discretized_gaussian_log_likelihood( |
| 434 | x_start, means=out["mean"], log_scales=0.5 * out["log_variance"] |
| 435 | ) |
| 436 | assert decoder_nll.shape == x_start.shape |
| 437 | decoder_nll = mean_flat(decoder_nll) / np.log(2.0) |
| 438 | |
| 439 | # At the first timestep return the decoder NLL, |
| 440 | # otherwise return KL(q(x_{t-1}|x_t,x_0) || p(x_{t-1}|x_t)) |
| 441 | output = torch.where((t == 0), decoder_nll, kl) |
| 442 | return {"output": output, "pred_xstart": out["pred_xstart"]} |
| 443 | |
| 444 | def calc_bpd_loop(self, model, x_start, clip_denoised=False, model_kwargs=None): |
| 445 | """ |
no test coverage detected