| 30 | del self.decoder.encoder |
| 31 | |
| 32 | def forward(self, latent_inputs, latent_predictions, image_inputs, split="train"): |
| 33 | log = dict() |
| 34 | loss = (latent_inputs - latent_predictions) ** 2 |
| 35 | log[f"{split}/latent_l2_loss"] = loss.mean().detach() |
| 36 | image_reconstructions = None |
| 37 | if self.perceptual_weight > 0.0: |
| 38 | image_reconstructions = self.decoder.decode(latent_predictions) |
| 39 | image_targets = self.decoder.decode(latent_inputs) |
| 40 | perceptual_loss = self.perceptual_loss(image_targets.contiguous(), image_reconstructions.contiguous()) |
| 41 | loss = self.latent_weight * loss.mean() + self.perceptual_weight * perceptual_loss.mean() |
| 42 | log[f"{split}/perceptual_loss"] = perceptual_loss.mean().detach() |
| 43 | |
| 44 | if self.perceptual_weight_on_inputs > 0.0: |
| 45 | image_reconstructions = default(image_reconstructions, self.decoder.decode(latent_predictions)) |
| 46 | if self.scale_input_to_tgt_size: |
| 47 | image_inputs = torch.nn.functional.interpolate( |
| 48 | image_inputs, |
| 49 | image_reconstructions.shape[2:], |
| 50 | mode="bicubic", |
| 51 | antialias=True, |
| 52 | ) |
| 53 | elif self.scale_tgt_to_input_size: |
| 54 | image_reconstructions = torch.nn.functional.interpolate( |
| 55 | image_reconstructions, |
| 56 | image_inputs.shape[2:], |
| 57 | mode="bicubic", |
| 58 | antialias=True, |
| 59 | ) |
| 60 | |
| 61 | perceptual_loss2 = self.perceptual_loss(image_inputs.contiguous(), image_reconstructions.contiguous()) |
| 62 | loss = loss + self.perceptual_weight_on_inputs * perceptual_loss2.mean() |
| 63 | log[f"{split}/perceptual_loss_on_inputs"] = perceptual_loss2.mean().detach() |
| 64 | return loss, log |