Compute training losses for a single timestep. :param model: the model to evaluate loss on. :param x_start: the [N x C x ...] tensor of inputs. :param t: a batch of timestep indices. :param model_kwargs: if not None, a dict of extra keyword arguments to
(
self, model, x_start, t, model_kwargs=None, noise=None
)
| 760 | } |
| 761 | |
| 762 | def training_losses( |
| 763 | self, model, x_start, t, model_kwargs=None, noise=None |
| 764 | ) -> Dict[str, th.Tensor]: |
| 765 | """ |
| 766 | Compute training losses for a single timestep. |
| 767 | |
| 768 | :param model: the model to evaluate loss on. |
| 769 | :param x_start: the [N x C x ...] tensor of inputs. |
| 770 | :param t: a batch of timestep indices. |
| 771 | :param model_kwargs: if not None, a dict of extra keyword arguments to |
| 772 | pass to the model. This can be used for conditioning. |
| 773 | :param noise: if specified, the specific Gaussian noise to try to remove. |
| 774 | :return: a dict with the key "loss" containing a tensor of shape [N]. |
| 775 | Some mean or variance settings may also have other keys. |
| 776 | """ |
| 777 | x_start = self.scale_channels(x_start) |
| 778 | if model_kwargs is None: |
| 779 | model_kwargs = {} |
| 780 | if noise is None: |
| 781 | noise = th.randn_like(x_start) |
| 782 | x_t = self.q_sample(x_start, t, noise=noise) |
| 783 | |
| 784 | terms = {} |
| 785 | |
| 786 | if self.loss_type == "kl" or self.loss_type == "rescaled_kl": |
| 787 | vb_terms = self._vb_terms_bpd( |
| 788 | model=model, |
| 789 | x_start=x_start, |
| 790 | x_t=x_t, |
| 791 | t=t, |
| 792 | clip_denoised=False, |
| 793 | model_kwargs=model_kwargs, |
| 794 | ) |
| 795 | terms["loss"] = vb_terms["output"] |
| 796 | if self.loss_type == "rescaled_kl": |
| 797 | terms["loss"] *= self.num_timesteps |
| 798 | extra = vb_terms["extra"] |
| 799 | elif self.loss_type == "mse" or self.loss_type == "rescaled_mse": |
| 800 | model_output = model(x_t, t, **model_kwargs) |
| 801 | if isinstance(model_output, tuple): |
| 802 | model_output, extra = model_output |
| 803 | else: |
| 804 | extra = {} |
| 805 | |
| 806 | if self.model_var_type in [ |
| 807 | "learned", |
| 808 | "learned_range", |
| 809 | ]: |
| 810 | B, C = x_t.shape[:2] |
| 811 | assert model_output.shape == (B, C * 2, *x_t.shape[2:]) |
| 812 | model_output, model_var_values = th.split(model_output, C, dim=1) |
| 813 | # Learn the variance using the variational bound, but don't let |
| 814 | # it affect our mean prediction. |
| 815 | frozen_out = th.cat([model_output.detach(), model_var_values], dim=1) |
| 816 | terms["vb"] = self._vb_terms_bpd( |
| 817 | model=lambda *args, r=frozen_out: r, |
| 818 | x_start=x_start, |
| 819 | x_t=x_t, |
nothing calls this directly
no test coverage detected