cc

.gitignore

@@ -1,4 +1,5 @@
 __pycache__
 datasets
 wandb
+checkpoints
 *.zip

baselines/cyclegan-cut/train.py

@@ -14,6 +14,7 @@ if __name__ == "__main__":
     dataset_size = len(dataset)  # get the number of images in the dataset.
 
     model = create_model(opt)  # create a model given opt.model and other options
+
     print("The number of training images = %d" % dataset_size)
 
     visualizer = Visualizer(

main.py

@@ -6,12 +6,16 @@ from swim.modules.dataset import SwimDataModule
 from lightning import Trainer
 from lightning.pytorch.loggers import WandbLogger
 
+torch.set_float32_matmul_precision("medium")
+
 config = OmegaConf.load("configs/autoencoder/autoencoder_kl_32x32x4.yaml")
 
 model = instantiate_from_config(config.model)
 model.learning_rate = config.model.base_learning_rate
 
-datamodule = SwimDataModule(img_size=32)
+datamodule = SwimDataModule(
+    root_dir="/cm/shared/ninhnq3/datasets/swim_data", batch_size=2, img_size=512
+)
 
 logger = WandbLogger(project="swim", name="autoencoder_kl")
 

swim/modules/losses/contperceptual.py

@@ -94,24 +94,31 @@ class LPIPSWithDiscriminator(nn.Module):
         split="train",
         weights=None,
     ):
-        rec_loss = torch.abs(inputs.contiguous() - reconstructions.contiguous())
-        if self.perceptual_weight > 0:
-            p_loss = self.perceptual_loss(
-                inputs.contiguous(), reconstructions.contiguous()
-            )
-            rec_loss = rec_loss + self.perceptual_weight * p_loss
-
-        nll_loss = rec_loss / torch.exp(self.logvar) + self.logvar
-        weighted_nll_loss = nll_loss
-        if weights is not None:
-            weighted_nll_loss = weights * nll_loss
-        weighted_nll_loss = torch.sum(weighted_nll_loss) / weighted_nll_loss.shape[0]
-        nll_loss = torch.sum(nll_loss) / nll_loss.shape[0]
-        kl_loss = posteriors.kl()
-        kl_loss = torch.sum(kl_loss) / kl_loss.shape[0]
 
         # now the GAN part
         if optimizer_idx == 0:
+
+            rec_loss = torch.abs(inputs.contiguous() - reconstructions.contiguous())
+
+            if self.perceptual_weight > 0:
+                p_loss = self.perceptual_loss(
+                    inputs.contiguous(), reconstructions.contiguous()
+                )
+                rec_loss = rec_loss + self.perceptual_weight * p_loss
+
+            nll_loss = rec_loss / torch.exp(self.logvar) + self.logvar
+            weighted_nll_loss = nll_loss
+            if weights is not None:
+                weighted_nll_loss = weights * nll_loss
+            weighted_nll_loss = (
+                torch.sum(weighted_nll_loss) / weighted_nll_loss.shape[0]
+            )
+            print(nll_loss.shape)
+
+            nll_loss = torch.sum(nll_loss) / nll_loss.shape[0]
+            kl_loss = posteriors.kl()
+            kl_loss = torch.sum(kl_loss) / kl_loss.shape[0]
+
             # generator update
             if cond is None:
                 assert not self.disc_conditional