removed commented unused code and edited annotations

Author: codinglabsong

README.md

@@ -1,5 +1,3 @@
-added spectral norm to path gan.
-learning reduces to 0.1
 different lr for gen and disc
 
 

src/aging_gan/data.py

@@ -46,7 +46,7 @@ def __getitem__(self, idx: int) -> Tuple[Tensor, int]:
 
 def make_unpaired_loader(
     root: str,
-    split: str,  # "train" | "valid" | "test"
+    split: str,
     transform: T.Compose,
     batch_size: int = 4,
     num_workers: int = 1,

src/aging_gan/model.py

@@ -4,9 +4,6 @@
 import torch.nn as nn
 import torch.nn.functional as F
 
-# import torch.nn.utils as nn_utils
-# import segmentation_models_pytorch as smp
-
 
 class ResidualBlock(nn.Module):
     """Simple residual block with two conv layers."""
@@ -138,56 +135,12 @@ def __init__(self, ndf: int) -> None:
 
     def forward(self, x: Tensor) -> Tensor:
         """Return discriminator logits for input ``x``."""
-        x = self.model(x)
-        return F.avg_pool2d(x, x.size()[2:]).view(x.size()[0], -1)
-
-
-# # Discriminator: PatchGAN 70x70
-# class PatchDiscriminator(nn.Module):
-#     def __init__(self, in_channels=3, ndf=48):
-#         super().__init__()
-#         layers = [
-#             nn_utils.spectral_norm(
-#                 nn.Conv2d(
-#                     in_channels=in_channels,
-#                     out_channels=ndf,
-#                     kernel_size=4,
-#                     stride=2,
-#                     padding=1,
-#                 )
-#             ),
-#             nn.LeakyReLU(0.2),
-#         ]
-#         nf = ndf
-#         for i in range(3):
-#             stride = 2 if i < 2 else 1
-#             layers += [
-#                 nn_utils.spectral_norm(nn.Conv2d(nf, nf * 2, 4, stride, 1)),
-#                 nn.InstanceNorm2d(nf * 2, affine=True),
-#                 nn.LeakyReLU(0.2),
-#             ]
-#             nf *= 2
-#         layers += [nn_utils.spectral_norm(nn.Conv2d(nf, 1, 4, 1, 1))]
-#         self.model = nn.Sequential(*layers)
-
-#     def forward(self, x):
-#         return self.model(x)
-
-
-# # Freeze encoder of model so that model can learn "aging" during the first epoch
-# def freeze_encoders(G, F):
-#     for param in G.encoder.parameters():
-#         param.requires_grad = False
-#     for param in F.encoder.parameters():
-#         param.requires_grad = False
-
-
-# # Unfreeze encoders later
-# def unfreeze_encoders(G, F):
-#     for param in G.encoder.parameters():
-#         param.requires_grad = True
-#     for param in F.encoder.parameters():
-#         param.requires_grad = True
+        # x: (B, 3, H, W)
+        x = self.model(x)  # (B, 1, H//8-2, W//8-2)
+        # Average pooling and flatten
+        return F.avg_pool2d(x, x.size()[2:]).view(
+            x.size()[0], -1
+        )  # global average -> (B, 1, 1, 1) -> flatten to (B, 1)
 
 
 # Initialize and return the generators and discriminators used for training
@@ -197,20 +150,6 @@ def initialize_models(
     n_blocks: int = 9,
 ) -> tuple[Generator, Generator, Discriminator, Discriminator]:
     """Instantiate generators and discriminators with default sizes."""
-    # G = smp.Unet(
-    #     encoder_name="resnet34",
-    #     encoder_weights="imagenet",  # preload low-level filters
-    #     in_channels=3,  # RGB input
-    #     classes=3,  # RGB output
-    # )
-
-    # F = smp.Unet(
-    #     encoder_name="resnet34",
-    #     encoder_weights="imagenet",  # preload low-level filters
-    #     in_channels=3,  # RGB input
-    #     classes=3,  # RGB output
-    # )
-
     # initialize the generators and discriminators
     G = Generator(ngf, n_blocks)
     F = Generator(ngf, n_blocks)

src/aging_gan/train.py

@@ -554,13 +554,6 @@ def main() -> None:
     # ---------- Models, Optimizers, Loss Functions, Schedulers Initialization ----------
     # Initialize the generators (G, F) and discriminators (DX, DY)
     G, F, DX, DY = initialize_models()
-    # Freeze generator encoderes for training during early epochs
-    # logger.info("Parameters of generator G:")
-    # logger.info(print_trainable_parameters(G))
-    # logger.info("Freezing encoders of generators...")
-    # freeze_encoders(G, F)
-    # logger.info("Parameters of generator G after freezing:")
-    # logger.info(print_trainable_parameters(G))
     # Initialize optimizers
     (
         opt_G,
@@ -613,13 +606,6 @@ def main() -> None:
     best_fid = float("inf")  # keep track of the best FID score for each epoch
     for epoch in range(1, cfg.num_train_epochs + 1):
         logger.info(f"\nEPOCH {epoch}")
-        # after 1 full epoch, unfreeze
-        # if epoch == 2:
-        #     logger.info("Unfreezing encoders of generators...")
-        #     unfreeze_encoders(G, F)
-        #     logger.info("Parameters of generator G after unfreezing:")
-        #     logger.info(print_trainable_parameters(G))
-
         val_metrics = perform_epoch(
             cfg,
             train_loader,

src/aging_gan/utils.py

@@ -39,20 +39,6 @@ def load_environ_vars(wandb_project: str = "aging-gan") -> None:
     logger.info(f"W&B project set to '{wandb_project}'")
 
 
-# def print_trainable_parameters(model) -> str:
-#     """
-#     Compute and return a summary of trainable vs. total parameters in a model.
-#     """
-#     trainable_params = 0
-#     all_param = 0
-#     for _, param in model.named_parameters():
-#         all_param += param.numel()
-#         if param.requires_grad:
-#             trainable_params += param.numel()
-
-#     return f"trainable params: {trainable_params} || all params: {all_param} || trainable%: {100 * trainable_params / all_param:.2f}"
-
-
 def save_checkpoint(
     epoch,
     G,