| 387 | return loss |
| 388 | |
| 389 | def p_losses(self, x_start, t, noise=None): |
| 390 | noise = default(noise, lambda: torch.randn_like(x_start)) |
| 391 | x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise) |
| 392 | model_out = self.model(x_noisy, t) |
| 393 | |
| 394 | loss_dict = {} |
| 395 | if self.parameterization == "eps": |
| 396 | target = noise |
| 397 | elif self.parameterization == "x0": |
| 398 | target = x_start |
| 399 | elif self.parameterization == "v": |
| 400 | target = self.get_v(x_start, noise, t) |
| 401 | else: |
| 402 | raise NotImplementedError(f"Parameterization {self.parameterization} not yet supported") |
| 403 | |
| 404 | loss = self.get_loss(model_out, target, mean=False).mean(dim=[1, 2, 3]) |
| 405 | |
| 406 | log_prefix = 'train' if self.training else 'val' |
| 407 | |
| 408 | loss_dict.update({f'{log_prefix}/loss_simple': loss.mean()}) |
| 409 | loss_simple = loss.mean() * self.l_simple_weight |
| 410 | |
| 411 | loss_vlb = (self.lvlb_weights[t] * loss).mean() |
| 412 | loss_dict.update({f'{log_prefix}/loss_vlb': loss_vlb}) |
| 413 | |
| 414 | loss = loss_simple + self.original_elbo_weight * loss_vlb |
| 415 | |
| 416 | loss_dict.update({f'{log_prefix}/loss': loss}) |
| 417 | |
| 418 | return loss, loss_dict |
| 419 | |
| 420 | def forward(self, x, *args, **kwargs): |
| 421 | # b, c, h, w, device, img_size, = *x.shape, x.device, self.image_size |