Compute the entire variational lower-bound, measured in bits-per-dim, as well as other related quantities. :param model: the model to evaluate loss on. :param x_start: the [N x C x ...] tensor of inputs. :param clip_denoised: if True, clip denoised samples.
(self, model, x_start, clip_denoised=False, model_kwargs=None)
| 863 | return mean_flat(kl_prior) / np.log(2.0) |
| 864 | |
| 865 | def calc_bpd_loop(self, model, x_start, clip_denoised=False, model_kwargs=None): |
| 866 | """ |
| 867 | Compute the entire variational lower-bound, measured in bits-per-dim, |
| 868 | as well as other related quantities. |
| 869 | |
| 870 | :param model: the model to evaluate loss on. |
| 871 | :param x_start: the [N x C x ...] tensor of inputs. |
| 872 | :param clip_denoised: if True, clip denoised samples. |
| 873 | :param model_kwargs: if not None, a dict of extra keyword arguments to |
| 874 | pass to the model. This can be used for conditioning. |
| 875 | |
| 876 | :return: a dict containing the following keys: |
| 877 | - total_bpd: the total variational lower-bound, per batch element. |
| 878 | - prior_bpd: the prior term in the lower-bound. |
| 879 | - vb: an [N x T] tensor of terms in the lower-bound. |
| 880 | - xstart_mse: an [N x T] tensor of x_0 MSEs for each timestep. |
| 881 | - mse: an [N x T] tensor of epsilon MSEs for each timestep. |
| 882 | """ |
| 883 | device = x_start.device |
| 884 | batch_size = x_start.shape[0] |
| 885 | |
| 886 | vb = [] |
| 887 | xstart_mse = [] |
| 888 | mse = [] |
| 889 | for t in list(range(self.num_timesteps))[::-1]: |
| 890 | t_batch = th.tensor([t] * batch_size, device=device) |
| 891 | noise = th.randn_like(x_start) |
| 892 | x_t = self.q_sample(x_start=x_start, t=t_batch, noise=noise) |
| 893 | # Calculate VLB term at the current timestep |
| 894 | with th.no_grad(): |
| 895 | out = self._vb_terms_bpd( |
| 896 | model, |
| 897 | x_start=x_start, |
| 898 | x_t=x_t, |
| 899 | t=t_batch, |
| 900 | clip_denoised=clip_denoised, |
| 901 | model_kwargs=model_kwargs, |
| 902 | ) |
| 903 | vb.append(out["output"]) |
| 904 | xstart_mse.append(mean_flat((out["pred_xstart"] - x_start) ** 2)) |
| 905 | eps = self._predict_eps_from_xstart(x_t, t_batch, out["pred_xstart"]) |
| 906 | mse.append(mean_flat((eps - noise) ** 2)) |
| 907 | |
| 908 | vb = th.stack(vb, dim=1) |
| 909 | xstart_mse = th.stack(xstart_mse, dim=1) |
| 910 | mse = th.stack(mse, dim=1) |
| 911 | |
| 912 | prior_bpd = self._prior_bpd(x_start) |
| 913 | total_bpd = vb.sum(dim=1) + prior_bpd |
| 914 | return { |
| 915 | "total_bpd": total_bpd, |
| 916 | "prior_bpd": prior_bpd, |
| 917 | "vb": vb, |
| 918 | "xstart_mse": xstart_mse, |
| 919 | "mse": mse, |
| 920 | } |
| 921 | |
| 922 | def scale_channels(self, x: th.Tensor) -> th.Tensor: |
nothing calls this directly
no test coverage detected