| 891 | return mean_flat(kl_prior) / np.log(2.0) |
| 892 | |
| 893 | def p_losses(self, x_start, cond, t, noise=None): |
| 894 | noise = default(noise, lambda: torch.randn_like(x_start)) |
| 895 | x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise) |
| 896 | model_output = self.apply_model(x_noisy, t, cond) |
| 897 | |
| 898 | loss_dict = {} |
| 899 | prefix = 'train' if self.training else 'val' |
| 900 | |
| 901 | if self.parameterization == "x0": |
| 902 | target = x_start |
| 903 | elif self.parameterization == "eps": |
| 904 | target = noise |
| 905 | elif self.parameterization == "v": |
| 906 | target = self.get_v(x_start, noise, t) |
| 907 | else: |
| 908 | raise NotImplementedError() |
| 909 | |
| 910 | loss_simple = self.get_loss(model_output, target, mean=False).mean([1, 2, 3]) |
| 911 | loss_dict.update({f'{prefix}/loss_simple': loss_simple.mean()}) |
| 912 | |
| 913 | logvar_t = self.logvar[t].to(self.device) |
| 914 | loss = loss_simple / torch.exp(logvar_t) + logvar_t |
| 915 | # loss = loss_simple / torch.exp(self.logvar) + self.logvar |
| 916 | if self.learn_logvar: |
| 917 | loss_dict.update({f'{prefix}/loss_gamma': loss.mean()}) |
| 918 | loss_dict.update({'logvar': self.logvar.data.mean()}) |
| 919 | |
| 920 | loss = self.l_simple_weight * loss.mean() |
| 921 | |
| 922 | loss_vlb = self.get_loss(model_output, target, mean=False).mean(dim=(1, 2, 3)) |
| 923 | loss_vlb = (self.lvlb_weights[t] * loss_vlb).mean() |
| 924 | loss_dict.update({f'{prefix}/loss_vlb': loss_vlb}) |
| 925 | loss += (self.original_elbo_weight * loss_vlb) |
| 926 | loss_dict.update({f'{prefix}/loss': loss}) |
| 927 | |
| 928 | return loss, loss_dict |
| 929 | |
| 930 | def p_mean_variance(self, x, c, t, clip_denoised: bool, return_codebook_ids=False, quantize_denoised=False, |
| 931 | return_x0=False, score_corrector=None, corrector_kwargs=None): |