File size: 3,840 Bytes
a04730e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 | from __future__ import annotations
from dataclasses import dataclass
import torch
from torch import nn
import torch.nn.functional as F
@dataclass
class VAELossOutput:
total_loss: torch.Tensor
recon_loss: torch.Tensor
kl_loss: torch.Tensor
perceptual_loss: torch.Tensor
class VAELoss(nn.Module):
"""
VAE loss:
total =
recon_weight * reconstruction_loss
+ kl_weight * KL
+ perceptual_weight * LPIPS
Inputs are expected to be in [-1, 1]
"""
def __init__(
self,
recon_loss_type: str = "l1",
recon_weight: float = 1.0,
kl_weight: float = 1e-6,
perceptual_weight: float = 0.0,
use_lpips: bool = False,
lpips_net: str = "vgg",
):
super().__init__()
if recon_loss_type not in {"l1", "mse"}:
raise ValueError(
f"Unknown recon_loss_type={recon_loss_type}. "
"Use 'l1' or 'mse'."
)
self.recon_loss_type = recon_loss_type
self.recon_weight = recon_weight
self.kl_weight = kl_weight
self.perceptual_weight = perceptual_weight
self.use_lpips = use_lpips
self.lpips_model = None
if use_lpips:
try:
import lpips
except ImportError as exc:
raise ImportError(
"LPIPS is enabled but package 'lpips' is not installed. "
"Install it with: pip install lpips"
) from exc
self.lpips_model = lpips.LPIPS(net=lpips_net)
self.lpips_model.eval()
for p in self.lpips_model.parameters():
p.requires_grad = False
def reconstruction_loss(
self,
x_recon: torch.Tensor,
x: torch.Tensor,
) -> torch.Tensor:
if self.recon_loss_type == "l1":
return F.l1_loss(x_recon, x)
if self.recon_loss_type == "mse":
return F.mse_loss(x_recon, x)
raise RuntimeError("Invalid reconstruction loss type.")
def perceptual_loss(
self,
x_recon: torch.Tensor,
x: torch.Tensor,
) -> torch.Tensor:
if not self.use_lpips or self.lpips_model is None:
return torch.zeros((), device=x.device, dtype=x.dtype)
# LPIPS expects images in [-1, 1], which matches our transform.
with torch.cuda.amp.autocast(enabled=False):
loss = self.lpips_model(
x_recon.float(),
x.float(),
).mean()
return loss.to(dtype=x.dtype)
def forward(
self,
x_recon: torch.Tensor,
x: torch.Tensor,
posterior,
kl_weight: float | None = None,
) -> VAELossOutput:
"""
Args:
x_recon:
Reconstructed image [B, 3, H, W], in [-1, 1].
x:
Target image [B, 3, H, W], in [-1, 1].
posterior:
DiagonalGaussianDistribution from vae.encode(x).
kl_weight:
Optional current KL weight. Useful for KL warmup.
Returns:
VAELossOutput.
"""
current_kl_weight = self.kl_weight if kl_weight is None else kl_weight
recon = self.reconstruction_loss(x_recon, x)
# posterior.kl() returns [B], already summed over latent dimensions.
kl = posterior.kl().mean()
perceptual = self.perceptual_loss(x_recon, x)
total = (
self.recon_weight * recon
+ current_kl_weight * kl
+ self.perceptual_weight * perceptual
)
return VAELossOutput(
total_loss=total,
recon_loss=recon.detach(),
kl_loss=kl.detach(),
perceptual_loss=perceptual.detach(),
) |