changed age limit and limit function for data preprocessing; removed not needed functions in utils and train

Author: codinglabsong

src/aging_gan/data.py

@@ -12,13 +12,13 @@
 
 class UTKFace(Dataset):
     """
-    Assumes the unzipped UTKFace images live in  <root>/data/utkface_aligned_cropped/UTKFace
+    Assumes the unzipped UTKFace images live in  <root>/data/UTKFace
     File pattern:  {age}_{gender}_{race}_{yyyymmddHHMMSS}.jpg
     """
 
     def __init__(self, root: str, transform: T.Compose | None = None):
         self.root = (
-            Path(root) / "utkface_aligned_cropped" / "UTKFace"
+            Path(root) / "UTKFace"  # "utkface_aligned_cropped" /
         )  # or "UTKFace" for the unaligned and varied original version.
         self.files = sorted(f for f in self.root.glob("*.jpg"))
         if not self.files:
@@ -46,10 +46,9 @@ def make_unpaired_loader(
     transform: T.Compose,
     batch_size: int = 4,
     num_workers: int = 1,
-    limit: int | None = None,  # per-domain cap
     seed: int = 42,
-    young_max: int = 25,  # 0-25
-    old_min: int = 55,  # 55+
+    young_max: int = 28,  # 18-28
+    old_min: int = 40,  # 40+
 ):
     full_ds = UTKFace(root, transform)
 
@@ -60,7 +59,7 @@ def make_unpaired_loader(
 
     for i, f in enumerate(full_ds.files):
         age = int(f.name.split("_")[0])
-        if age <= young_max:
+        if age <= young_max and age >= 18:
             young_idx.append(i)
         elif age >= old_min:
             old_idx.append(i)
@@ -84,10 +83,10 @@ def split_indices(idxs: list[int]):
     part_y = split_indices(young_idx)[split].tolist()
     part_o = split_indices(old_idx)[split].tolist()
 
-    # Limit per domain
-    if limit is not None:
-        part_y = part_y[:limit]
-        part_o = part_o[:limit]
+    # same dataset length
+    limit = min(len(part_y), len(part_o))
+    part_y = part_y[:limit]
+    part_o = part_o[:limit]
 
     # Wrap subsets in unpaird Dataset
     @dataclass
@@ -107,9 +106,7 @@ def __getitem__(self, idx: int):
     old_ds = Subset(full_ds, part_o)
     paired = Unpaired(young_ds, old_ds)
 
-    logger.info(
-        f"- UTK {split}: young={len(young_ds)}  old={len(old_ds)}" f"(limit={limit})"
-    )
+    logger.info(f"- UTK {split}: young={len(young_ds)}  old={len(old_ds)}")
     return DataLoader(
         paired,
         batch_size=batch_size,
@@ -126,11 +123,8 @@ def prepare_dataset(
     train_batch_size: int = 4,
     eval_batch_size: int = 8,
     num_workers: int = 2,
-    center_crop_size: int = 256,
+    img_size: int = 256,
     resize_size: int = 286,
-    train_size: int | None = None,  # None = use all
-    val_size: int | None = None,
-    test_size: int | None = None,
     seed: int = 42,
 ):
     data_dir = Path(__file__).resolve().parents[2] / "data"
@@ -139,19 +133,11 @@ def prepare_dataset(
     # randomness
     train_transform = T.Compose(
         [
-            T.Resize(resize_size, antialias=True),
-            T.CenterCrop(center_crop_size),
-            T.RandomApply(
-                [
-                    T.RandomAffine(
-                        degrees=5, translate=(0.02, 0.02), scale=(0.97, 1.03), shear=2,
-                        interpolation=T.InterpolationMode.BILINEAR, fill=0,
-                    )
-                ],
-                p=0.3,
-            ),
-            T.RandomHorizontalFlip(0.5),
-            T.ColorJitter(0.05, 0.05, 0.05, 0.02),
+            T.ToPILImage(),
+            T.RandomHorizontalFlip(),
+            T.Resize((img_size + 50, img_size + 50), antialias=True),
+            T.RandomCrop(img_size),
+            T.RandomRotation(degrees=(0, 80)),
             T.ToTensor(),
             T.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]),
         ]
@@ -160,8 +146,8 @@ def prepare_dataset(
     # deterministic
     eval_transform = T.Compose(
         [
-            T.CenterCrop(center_crop_size),
-            T.Resize(resize_size),
+            T.Resize(resize_size, antialias=True),
+            T.CenterCrop(img_size),
             T.ToTensor(),
             T.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]),
         ]
@@ -175,7 +161,6 @@ def prepare_dataset(
         train_transform,
         train_batch_size,
         num_workers,
-        train_size,
         seed,
     )
     val_loader = make_unpaired_loader(
@@ -184,7 +169,6 @@ def prepare_dataset(
         eval_transform,
         eval_batch_size,
         num_workers,
-        val_size,
         seed,
     )
     test_loader = make_unpaired_loader(
@@ -193,7 +177,6 @@ def prepare_dataset(
         eval_transform,
         eval_batch_size,
         num_workers,
-        test_size,
         seed,
     )
     logger.info("Done.")

src/aging_gan/train.py

@@ -15,13 +15,12 @@
 from aging_gan.utils import (
     set_seed,
     load_environ_vars,
-    print_trainable_parameters,
     save_checkpoint,
     generate_and_save_samples,
     get_device,
 )
 from aging_gan.data import prepare_dataset
-from aging_gan.model import initialize_models, freeze_encoders, unfreeze_encoders
+from aging_gan.model import initialize_models
 from aging_gan.utils import archive_and_terminate
 
 logger = logging.getLogger(__name__)
@@ -59,6 +58,12 @@ def parse_args() -> argparse.Namespace:
         default=32,
         help="Batch size per device during evaluation.",
     )
+    p.add_argument(
+        "--lambda_adv_value",
+        type=int,
+        default=2,
+        help="Weight for adversarial loss",
+    )
     p.add_argument(
         "--lambda_cyc_value",
         type=int,
@@ -98,24 +103,6 @@ def parse_args() -> argparse.Namespace:
         default=10,
         help="The number of example generated images to save per epoch.",
     )
-    p.add_argument(
-        "--train_size",
-        type=int,
-        default=3000,
-        help="The size of train dataset to train on.",
-    )
-    p.add_argument(
-        "--val_size",
-        type=int,
-        default=800,
-        help="The size of validation dataset to evaluate.",
-    )
-    p.add_argument(
-        "--test_size",
-        type=int,
-        default=800,
-        help="The size of test dataset to evaluate.",
-    )
     p.add_argument(
         "--num_workers",
         type=int,
@@ -176,13 +163,16 @@ def initialize_optimizers(cfg, G, F, DX, DY):
     return opt_G, opt_F, opt_DX, opt_DY
 
 
-def initialize_loss_functions(lambda_cyc_value: int = 10, lambda_id_value: int = 5):
+def initialize_loss_functions(
+    lambda_adv_value: int = 2, lambda_cyc_value: int = 10, lambda_id_value: int = 5
+):
     mse = nn.MSELoss()
     l1 = nn.L1Loss()
+    lambda_adv = lambda_adv_value
     lambda_cyc = lambda_cyc_value
     lambda_id = lambda_id_value
 
-    return mse, l1, lambda_cyc, lambda_id
+    return mse, l1, lambda_adv, lambda_cyc, lambda_id
 
 
 def make_schedulers(cfg, opt_G, opt_F, opt_DX, opt_DY):
@@ -212,6 +202,7 @@ def perform_train_step(
     real_data,
     mse,
     l1,
+    lambda_adv,
     lambda_cyc,
     lambda_id,  # loss functions and loss params
     opt_G,
@@ -260,10 +251,10 @@ def perform_train_step(
     opt_F.zero_grad(set_to_none=True)
     # Loss 1: adversarial terms
     fake_test_logits = DX(fake_x)  # fake x logits
-    loss_f_adv = mse(fake_test_logits, torch.ones_like(fake_test_logits))
+    loss_f_adv = lambda_adv * mse(fake_test_logits, torch.ones_like(fake_test_logits))
 
     fake_test_logits = DY(fake_y)  # fake y logits
-    loss_g_adv = mse(fake_test_logits, torch.ones_like(fake_test_logits))
+    loss_g_adv = lambda_adv * mse(fake_test_logits, torch.ones_like(fake_test_logits))
     # Loss 2: cycle terms
     loss_cyc = lambda_cyc * (l1(rec_x, x) + l1(rec_y, y))
     # Loss 3: identity terms
@@ -299,6 +290,7 @@ def evaluate_epoch(
     split: str,  # either "val" or "test"
     mse,
     l1,
+    lambda_adv,
     lambda_cyc,
     lambda_id,  # loss functions and loss params
     fid_metric,
@@ -349,10 +341,14 @@ def evaluate_epoch(
             # ------ Evaluate Generators ------
             # Loss 1: adversarial terms
             fake_test_logits = DX(fake_x)  # fake x logits
-            loss_f_adv = mse(fake_test_logits, torch.ones_like(fake_test_logits))
+            loss_f_adv = lambda_adv * mse(
+                fake_test_logits, torch.ones_like(fake_test_logits)
+            )
 
             fake_test_logits = DY(fake_y)  # fake y logits
-            loss_g_adv = mse(fake_test_logits, torch.ones_like(fake_test_logits))
+            loss_g_adv = lambda_adv * mse(
+                fake_test_logits, torch.ones_like(fake_test_logits)
+            )
             # Loss 2: cycle terms
             loss_cyc = lambda_cyc * (l1(rec_x, x) + l1(rec_y, y))
             # Loss 3: identity terms
@@ -396,6 +392,7 @@ def perform_epoch(
     DY,
     mse,
     l1,
+    lambda_adv,
     lambda_cyc,
     lambda_id,
     opt_G,
@@ -427,6 +424,7 @@ def perform_epoch(
             real_data,
             mse,
             l1,
+            lambda_adv,
             lambda_cyc,
             lambda_id,  # loss functions and loss params
             opt_G,
@@ -469,6 +467,7 @@ def perform_epoch(
         "val",
         mse,
         l1,
+        lambda_adv,
         lambda_cyc,
         lambda_id,  # loss functions and loss params
         fid_metric,  # evaluation metric
@@ -527,22 +526,19 @@ def main() -> None:
         cfg.train_batch_size,
         cfg.eval_batch_size,
         cfg.num_workers,
-        train_size=cfg.train_size,
-        val_size=cfg.val_size,
-        test_size=cfg.test_size,
         seed=cfg.seed,
     )
 
     # ---------- 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))
+    # 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,
@@ -579,8 +575,8 @@ def main() -> None:
         test_loader,
     )
     # Loss functions and scalers
-    mse, l1, lambda_cyc, lambda_id = initialize_loss_functions(
-        cfg.lambda_cyc_value, cfg.lambda_id_value
+    mse, l1, lambda_adv, lambda_cyc, lambda_id = initialize_loss_functions(
+        cfg.lambda_adv_value, cfg.lambda_cyc_value, cfg.lambda_id_value
     )
     # Initialize schedulers (It it important this comes AFTER wrapping optimizers in accelerator)
     sched_G, sched_F, sched_DX, sched_DY = make_schedulers(
@@ -596,11 +592,11 @@ def main() -> None:
     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))
+        # 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,
@@ -612,6 +608,7 @@ def main() -> None:
             DY,
             mse,
             l1,
+            lambda_adv,
             lambda_cyc,
             lambda_id,
             opt_G,

src/aging_gan/utils.py

@@ -34,18 +34,18 @@ def load_environ_vars(wandb_project: str = "aging-gan"):
     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 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(