cc

configs/experiment/ch_64.yaml

@@ -10,15 +10,20 @@ seed: 42
 
 trainer:
   max_epochs: 100
+  # precision: 16-mixed
 
 model:
-  channels: 64
-  z_c_channels: 256
-  updown_channel_mults: [1, 2, 4]
-  n_enc_resnet_blocks: 4
-  n_dec_resnet_blocks: 6
+  channels: 128
+  z_c_channels: 512
+  updown_channel_mults: [1, 2, 2, 4]
+  n_enc_resnet_blocks: 8
+  n_dec_resnet_blocks: 8
   n_f_d_resnet_blocks: 4
-  learning_rate: 1e-4
+  learning_rate: 1e-5
+  weight_decay: 1e-2
+  beta_1: 0.9
+  beta_2: 0.999
 
 data:
   batch_size: 4
+

swim/data/swim_data.py

@@ -26,9 +26,14 @@ class SwimDataset(Dataset):
         if split == "train":
             self.transform = T.Compose(
                 [
-                    # T.Resize(self.img_size),
-                    # T.RandomCrop(self.img_size),
-                    T.RandomResizedCrop(self.img_size, scale=(0.5, 1.0)),
+                    T.Resize(self.img_size),
+                    T.RandomCrop(self.img_size),
+                    # T.RandomResizedCrop(
+                    #     self.img_size,
+                    #     scale=(0.5, 1.0),
+                    #     ratio=(1.0, 1.0),
+                    #     interpolation=T.InterpolationMode.LANCZOS,
+                    # ),
                     T.RandomHorizontalFlip(),
                     T.ToTensor(),
                     T.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),

swim/models/discriminator.py

@@ -73,17 +73,7 @@ class FeatureDiscriminator(nn.Module):
         for _ in range(n_resnet_blocks):
             self.resnet_blocks.append(SNResnetBlock(channels, channels, d_style_emb))
 
-        self.conv_out = spectral_norm(
-            nn.Conv2d(channels, z_c_channels, 3, stride=1, padding=1)
-        )
-
-        self.mlp = nn.Sequential(
-            spectral_norm(nn.Linear(z_c_channels + d_style_emb, 256)),
-            nn.LeakyReLU(0.2),
-            spectral_norm(nn.Linear(256, 128)),
-            nn.LeakyReLU(0.2),
-            spectral_norm(nn.Linear(128, 1)),
-        )
+        self.conv_out = spectral_norm(nn.Conv2d(channels, 1, 3, stride=1, padding=1))
 
     def forward(
         self, x: torch.Tensor, style_emb: torch.Tensor, for_G=False, for_real=False
@@ -96,25 +86,25 @@ class FeatureDiscriminator(nn.Module):
             h = F.avg_pool2d(h, 2)
 
         h = self.conv_out(h)
-        h = F.leaky_relu(h, 0.2)
 
-        h = F.adaptive_avg_pool2d(h, 1)
-        h = torch.flatten(h, 1)
-        h = self.mlp(torch.concat([h, style_emb], dim=1))
+        if for_G:
+            for_real = True
 
-        # if for_real:
-        #     target = torch.full_like(h, 1.0)
-        # else:
-        #     target = torch.zeros_like(h)
+        if for_real:
+            target = torch.full_like(h, 1.0)
+        else:
+            target = torch.full_like(h, 0.0)
 
-        # loss = F.binary_cross_entropy_with_logits(h, target, reduction="none")
+        loss = F.binary_cross_entropy_with_logits(h, target, reduction="none").mean(
+            dim=[2, 3]
+        )
 
         # hinge loss
-        if for_G:
-            loss = -h
-        elif for_real:
-            loss = F.relu(1.0 - h)
-        else:
-            loss = F.relu(1.0 + h)
+        # if for_G:
+        #     loss = -h
+        # elif for_real:
+        #     loss = F.relu(1.0 - h)
+        # else:
+        #     loss = F.relu(1.0 + h)
 
         return loss

swim/models/style_encoder.py

@@ -11,11 +11,15 @@ class StyleEncoder(nn.Module):
         self.resnet = resnet18(weights=ResNet18_Weights.DEFAULT)
         self.resnet = nn.Sequential(*list(self.resnet.children())[:-1])
 
-        self.fc = nn.Linear(512, d_style_emb)
+        self.fc = nn.Sequential(
+            nn.Linear(512, 256),
+            nn.SiLU(),
+            nn.Linear(256, d_style_emb),
+        )
 
     def forward(self, x):
         # resize input to 224x224
-        # x = F.interpolate(x, size=(224, 224), mode="bilinear")
+        x = F.interpolate(x, size=(224, 224), mode="bilinear")
 
         x = self.resnet(x)
         x = torch.flatten(x, 1)

swim/models/swim_gan.py

@@ -9,7 +9,6 @@ from PIL import Image
 
 from lightning import LightningModule
 
-from diffusers import AutoencoderKL
 from diffusers.utils import make_image_grid
 
 from swim.utils.tensor_pool import GroupTensorPool
@@ -20,6 +19,7 @@ from .decoder import Decoder
 from .discriminator import FeatureDiscriminator
 
 import vision_aided_loss
+from lpips import LPIPS
 
 
 class SwimGAN(LightningModule):
@@ -29,21 +29,23 @@ class SwimGAN(LightningModule):
         channels: int = 128,
         z_c_channels: int = 512,
         updown_channel_mults: List[int] = [1, 2, 4],
-        n_enc_resnet_blocks: int = 4,
+        n_enc_resnet_blocks: int = 6,
         n_dec_resnet_blocks: int = 6,
         n_f_d_resnet_blocks: int = 2,
         n_styles: int = 5,
         d_style_emb: int = 128,
         input_size: int = 512,
         learning_rate: float = 1e-5,
-        weight_decay: float = 1e-4,
+        weight_decay: float = 1e-2,
+        beta_1: float = 0.9,
+        beta_2: float = 0.999,
         lambda_cls: float = 10.0,
-        lambda_rec: float = 1.0,
-        lambda_cycle: float = 1.0,
-        lambda_f_g: float = 10.0,
+        lambda_rec: float = 10.0,
+        lambda_cycle: float = 10.0,
+        lambda_f_g: float = 1.0,
         lambda_i_g: float = 1.0,
         lambda_c_const: float = 1.0,
-        lambda_s_const: float = 10.0,
+        lambda_s_const: float = 1.0,
     ):
         super().__init__()
 
@@ -53,6 +55,8 @@ class SwimGAN(LightningModule):
         self.n_styles = n_styles
         self.learning_rate = learning_rate
         self.weight_decay = weight_decay
+        self.beta_1 = beta_1
+        self.beta_2 = beta_2
         self.lambda_rec = lambda_rec
         self.lambda_cycle = lambda_cycle
         self.lambda_cls = lambda_cls
@@ -82,7 +86,10 @@ class SwimGAN(LightningModule):
 
         # training only
         self.i_discriminator = vision_aided_loss.Discriminator(
-            cv_type="clip", loss_type="multilevel_sigmoid", device="cpu"
+            cv_type="clip",
+            num_classes=n_styles,
+            loss_type="multilevel_sigmoid_s",
+            device="cpu",
         )
 
         self.f_discriminator = FeatureDiscriminator(
@@ -95,6 +102,9 @@ class SwimGAN(LightningModule):
         self.style_pool = GroupTensorPool(n_styles, 32)
         self.content_pool = GroupTensorPool(n_styles, 32)
 
+        # self.lpips = LPIPS(net="vgg")
+        # self.lpips.requires_grad_(False)
+
         self.cls_loss = nn.CrossEntropyLoss()
 
     def on_fit_start(self):
@@ -108,7 +118,7 @@ class SwimGAN(LightningModule):
 
         g_opt, i_d_opt, f_d_opt = self.optimizers()
 
-        # train the autoencoder
+        # train g1
         z_s = self.style_encoder(x)
         z_c = self.content_encoder(x)
         x_rec = self.decoder(z_c, z_s)
@@ -116,11 +126,28 @@ class SwimGAN(LightningModule):
         style_logits = self.style_classifier(z_s)
         cls_loss = self.cls_loss(style_logits, gt_style)
 
-        rec_loss = F.l1_loss(x, x_rec)
+        rec_loss = F.l1_loss(x, x_rec)  # + self.lpips(x, x_rec).mean()
+
+        f_g_loss = self.f_discriminator(z_c, z_s, for_G=True).mean()
+
+        # g1_loss = (
+        #     self.lambda_rec * rec_loss
+        #     + self.lambda_f_g * f_g_loss
+        #     + self.lambda_cls * cls_loss
+        # )
+
+        # g1_opt.zero_grad()
 
-        # # sample a random content and style feature
+        # self.manual_backward(g1_loss)
+        # g1_opt.step()
+
+        # sample a random content and style feature
         z_c_hat, _ = self.content_pool.query(z_c, gt_style)
-        z_s_hat, _ = self.style_pool.query(z_s, gt_style)
+        z_s_hat, gt_style_hat = self.style_pool.query(z_s, gt_style)
+
+        # train g2
+        z_c = z_c.detach()
+        z_s = z_s.detach()
 
         x1 = self.decoder(z_c, z_s_hat)
         x2 = self.decoder(z_c_hat, z_s)
@@ -128,29 +155,30 @@ class SwimGAN(LightningModule):
         z_c_rec = self.content_encoder(x1)
         z_s_rec = self.style_encoder(x2)
 
+        z_c_hat_rec = self.content_encoder(x2)
+        z_s_hat_rec = self.style_encoder(x1)
+
         x_cycle = self.decoder(z_c_rec, z_s_rec)
 
-        c_const_loss = F.l1_loss(z_c, z_c_rec)
-        s_const_loss = F.l1_loss(z_s, z_s_rec)
+        cycle_loss = F.l1_loss(x, x_cycle)  # + self.lpips(x, x_cycle).mean()
 
-        cycle_loss = F.l1_loss(x, x_cycle)
+        c_const_loss = F.l1_loss(z_c, z_c_rec) + F.l1_loss(z_c_hat, z_c_hat_rec)
+        s_const_loss = F.l1_loss(z_s, z_s_rec) + F.l1_loss(z_s_hat, z_s_hat_rec)
 
-        # adversarial loss
         i_g_loss = (
-            self.i_discriminator(x1, for_G=True).mean()
-            + self.i_discriminator(x2, for_G=True).mean()
-        ) / 2
-
-        c_g_loss = self.f_discriminator(z_c, z_s, for_G=True).mean()
+            self.i_discriminator(x1, gt_style_hat, for_G=True).mean()
+            + self.i_discriminator(x2, gt_style, for_G=True).mean()
+            + self.i_discriminator(x_cycle, gt_style, for_G=True).mean()
+        ) / 3.0
 
         g_loss = (
             self.lambda_rec * rec_loss
+            + self.lambda_f_g * f_g_loss
+            + self.lambda_cls * cls_loss
             + self.lambda_cycle * cycle_loss
             + self.lambda_c_const * c_const_loss
             + self.lambda_s_const * s_const_loss
             + self.lambda_i_g * i_g_loss
-            + self.lambda_f_g * c_g_loss
-            + self.lambda_cls * cls_loss
         )
 
         g_opt.zero_grad()
@@ -158,14 +186,19 @@ class SwimGAN(LightningModule):
         g_opt.step()
 
         # train the image discriminator
+        x1 = x1.detach()
+        x2 = x2.detach()
+        x_cycle = x_cycle.detach()
+
         i_d_loss = (
-            self.i_discriminator(x, for_real=True).mean()
+            self.i_discriminator(x, gt_style, for_real=True).mean()
             + (
-                self.i_discriminator(x1.detach(), for_real=False).mean()
-                + self.i_discriminator(x2.detach(), for_real=False).mean()
+                self.i_discriminator(x1, gt_style_hat, for_real=False).mean()
+                + self.i_discriminator(x2, gt_style, for_real=False).mean()
+                + self.i_discriminator(x_cycle, gt_style, for_real=False).mean()
             )
-            / 2
-        ) / 2
+            / 3.0
+        ) / 2.0
 
         i_d_opt.zero_grad()
         self.manual_backward(i_d_loss)
@@ -173,10 +206,10 @@ class SwimGAN(LightningModule):
 
         # train the feature discriminator
         f_d_loss = (
-            self.f_discriminator(z_c.detach(), z_s.detach(), for_real=False).mean()
+            self.f_discriminator(z_c, z_s, for_real=False).mean()
             + (
-                self.f_discriminator(z_c.detach(), z_s_hat, for_real=True).mean()
-                + self.f_discriminator(z_c_hat, z_s.detach(), for_real=True).mean()
+                self.f_discriminator(z_c, z_s_hat, for_real=True).mean()
+                + self.f_discriminator(z_c_hat, z_s, for_real=True).mean()
             )
             / 2
         ) / 2
@@ -191,8 +224,8 @@ class SwimGAN(LightningModule):
                 "train/cycle_loss": cycle_loss,
                 "train/cls_loss": cls_loss,
                 "train/i_g_loss": i_g_loss,
-                "train/i_d_loss": f_d_loss,
-                "train/f_g_loss": c_g_loss,
+                "train/i_d_loss": i_d_loss,
+                "train/f_g_loss": f_g_loss,
                 "train/f_d_loss": f_d_loss,
                 "train/c_const_loss": c_const_loss,
                 "train/s_const_loss": s_const_loss,
@@ -257,31 +290,38 @@ class SwimGAN(LightningModule):
         return x
 
     def configure_optimizers(self):
-        g_opt = torch.optim.AdamW(
+        g1_opt = torch.optim.AdamW(
             [
                 {"params": self.content_encoder.parameters()},
-                {"params": self.style_encoder.fc.parameters()},
-                {
-                    "params": self.style_encoder.resnet.parameters(),
-                    "lr": self.learning_rate / 10,
-                },
+                {"params": self.style_encoder.parameters()},
                 {"params": self.style_classifier.parameters()},
                 {"params": self.decoder.parameters()},
             ],
             lr=self.learning_rate,
             weight_decay=self.weight_decay,
+            betas=(self.beta_1, self.beta_2),
         )
 
+        # g2_opt = torch.optim.AdamW(
+        #     [
+        #         {"params": self.decoder.parameters()},
+        #     ],
+        #     lr=self.learning_rate,
+        #     weight_decay=self.weight_decay,
+        # )
+
         i_d_opt = torch.optim.AdamW(
             list(self.i_discriminator.parameters()),
             lr=self.learning_rate,
             weight_decay=self.weight_decay,
+            betas=(self.beta_1, self.beta_2),
         )
 
         f_d_opt = torch.optim.AdamW(
             list(self.f_discriminator.parameters()),
-            lr=self.learning_rate * 2,
+            lr=self.learning_rate * 10,
             weight_decay=self.weight_decay,
+            betas=(self.beta_1, self.beta_2),
         )
 
-        return [g_opt, i_d_opt, f_d_opt]
+        return [g1_opt, i_d_opt, f_d_opt]

swim/train.py

@@ -88,7 +88,7 @@ def train(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]:
 
     if cfg.get("train"):
         if cfg.compile:
-            model.compile()
+            model = torch.compile(model)
 
         log.info("Starting training!")
         trainer.fit(model=model, datamodule=datamodule, ckpt_path=cfg.get("ckpt_path"))

swim/utils/tensor_pool.py

@@ -28,6 +28,10 @@ class TensorPool(object):
         return torch.stack(return_tensors)
 
 
+import torch
+from random import choice, randint
+
+
 class GroupTensorPool(object):
 
     def __init__(self, n_groups: int = 5, pool_size: int = 32):
@@ -40,25 +44,34 @@ class GroupTensorPool(object):
         return_groups = []
 
         tensors = tensors.detach().clone()
+
         for tensor, curr_group in zip(tensors, labels):
-            # choose a random group except the current one
-            # if the pool is empty, return the current tensor
+            self.save_new_tensor(tensor, curr_group.item())
+
+        for tensor, curr_group in zip(tensors, labels):
+            # Choose a random group except the current one
             group = choice([i for i in range(self.n_groups) if i != curr_group.item()])
             pool = self.pools[group]
 
             if len(pool) == 0:
+                # If the selected group pool is empty, return the current tensor
                 return_tensors.append(tensor)
                 return_groups.append(curr_group)
-                pool.append((tensor, curr_group))
-            elif len(pool) < self.pool_size:
-                idx = randint(0, len(pool) - 1)
-                return_tensors.append(pool[idx][0])
-                return_groups.append(pool[idx][1])
-                pool.append((tensor, curr_group))
             else:
-                idx = randint(0, self.pool_size - 1)
-                return_tensors.append(pool[idx][0])
-                return_groups.append(pool[idx][1])
-                pool[idx] = (tensor, curr_group)
+                # Otherwise, select a random tensor from the pool
+                random_tensor = choice(pool)
+                return_tensors.append(random_tensor)
+                return_groups.append(torch.tensor(group))
 
         return torch.stack(return_tensors), torch.stack(return_groups)
+
+    def save_new_tensor(self, tensor: torch.Tensor, group: int):
+        pool = self.pools[group]
+
+        if len(pool) < self.pool_size:
+            # If the pool is not full, append the tensor
+            pool.append(tensor)
+        else:
+            # Replace a random item in the pool with the new tensor
+            replace_idx = randint(0, len(pool) - 1)
+            pool[replace_idx] = tensor

train_swim.sh

@@ -3,7 +3,7 @@ python swim/train.py \
     data.root_dir=/cm/shared/ninhnq3/datasets/swim_data \
     logger=wandb \
     logger.wandb.save_dir=/cm/shared/ninhnq3/workdirs/swim \
-    +trainer.val_check_interval=0.02 \
+    +trainer.val_check_interval=0.05 \
     +trainer.limit_val_batches=0.01 \
     callbacks.model_checkpoint.dirpath=/cm/shared/ninhnq3/checkpoints/swim_final