Training loss calculation
(self, denoise_fn, data_start, t, noise=None)
| 303 | return (kl, pred_xstart) if return_pred_xstart else kl |
| 304 | |
| 305 | def p_losses(self, denoise_fn, data_start, t, noise=None): |
| 306 | """ |
| 307 | Training loss calculation |
| 308 | """ |
| 309 | B, D, N = data_start.shape |
| 310 | assert t.shape == torch.Size([B]) |
| 311 | |
| 312 | if noise is None: |
| 313 | noise = torch.randn(data_start.shape, dtype=data_start.dtype, device=data_start.device) |
| 314 | assert noise.shape == data_start.shape and noise.dtype == data_start.dtype |
| 315 | |
| 316 | data_t = self.q_sample(x_start=data_start, t=t, noise=noise) |
| 317 | |
| 318 | if self.loss_type == 'mse': |
| 319 | # predict the noise instead of x_start. seems to be weighted naturally like SNR |
| 320 | eps_recon = denoise_fn(data_t, t) |
| 321 | assert data_t.shape == data_start.shape |
| 322 | assert eps_recon.shape == torch.Size([B, D, N]) |
| 323 | assert eps_recon.shape == data_start.shape |
| 324 | losses = ((noise - eps_recon)**2).mean(dim=list(range(1, len(data_start.shape)))) |
| 325 | elif self.loss_type == 'kl': |
| 326 | losses = self._vb_terms_bpd( |
| 327 | denoise_fn=denoise_fn, data_start=data_start, data_t=data_t, t=t, clip_denoised=False, |
| 328 | return_pred_xstart=False) |
| 329 | else: |
| 330 | raise NotImplementedError(self.loss_type) |
| 331 | |
| 332 | assert losses.shape == torch.Size([B]) |
| 333 | return losses |
| 334 | |
| 335 | '''debug''' |
| 336 |