cc

swim/models/autoencoder.py

@@ -10,7 +10,9 @@ from torchmetrics import (
     PeakSignalNoiseRatio,
     StructuralSimilarityIndexMeasure,
     MeanSquaredError,
+    MeanMetric,
 )
+from lpips import LPIPS
 
 
 class Autoencoder(LightningModule):
@@ -63,12 +65,15 @@ class Autoencoder(LightningModule):
         # embedding space
         self.post_quant_conv = nn.Conv2d(emb_channels, z_channels, 1)
 
+        self.lpips = LPIPS(net="vgg").eval()
+
         self.train_psnr = PeakSignalNoiseRatio()
         self.train_ssim = StructuralSimilarityIndexMeasure()
 
         self.val_psnr = PeakSignalNoiseRatio()
         self.val_ssim = StructuralSimilarityIndexMeasure()
         self.val_mse = MeanSquaredError()
+        self.val_lpips = MeanMetric()
 
     def encode(self, img: torch.Tensor) -> GaussianDistribution:
         """
@@ -114,17 +119,20 @@ class Autoencoder(LightningModule):
         img = batch["images"]
         recon = self.forward(img)
         # Calculate the loss
-        loss = torch.abs(img - recon).sum()  # L1 loss
+        l1_loss = torch.abs(img - recon).sum()  # L1 loss
+        lpips_loss = self.lpips.forward(recon, img).sum()  # LPIPS loss
+        total_loss = l1_loss + lpips_loss
 
         self.train_psnr(recon, img)
         self.train_ssim(recon, img)
 
         # Log the loss
-        self.log("train/l1_loss", loss.item(), on_step=True, prog_bar=True)
+        self.log("train/l1_loss", l1_loss.item(), on_step=True, prog_bar=True)
+        self.log("train/lpips_loss", lpips_loss.item(), on_step=True, prog_bar=True)
         self.log("train/psnr", self.train_psnr, on_step=True, prog_bar=True)
         self.log("train/ssim", self.train_ssim, on_step=True, prog_bar=True)
 
-        return loss
+        return total_loss
 
     def validation_step(self, batch, batch_idx):
         """
@@ -138,14 +146,20 @@ class Autoencoder(LightningModule):
         # Get the distribution
         recon = self.forward(img)
 
+        lpips_loss = self.lpips.forward(recon, img)  # LPIPS loss
+
         self.val_psnr(recon, img)
         self.val_ssim(recon, img)
         self.val_mse(recon, img)
+        self.val_lpips(lpips_loss)
 
         # Log the loss
         self.log("val/psnr", self.val_psnr, on_epoch=True, on_step=False, prog_bar=True)
         self.log("val/ssim", self.val_ssim, on_epoch=True, on_step=False, prog_bar=True)
         self.log("val/mse", self.val_mse, on_epoch=True, on_step=False, prog_bar=True)
+        self.log(
+            "val/lpips", self.val_lpips, on_epoch=True, on_step=False, prog_bar=True
+        )
 
         if batch_idx == 0:
             self.log_images(img, recon)