tuned hyperparams; edited save checkpoint to every 5 batches

Author: codinglabsong

src/aging_gan/data.py

@@ -124,7 +124,6 @@ def prepare_dataset(
     eval_batch_size: int = 8,
     num_workers: int = 2,
     img_size: int = 256,
-    resize_size: int = 286,
     seed: int = 42,
 ):
     data_dir = Path(__file__).resolve().parents[2] / "data"
@@ -146,7 +145,7 @@ def prepare_dataset(
     # deterministic
     eval_transform = T.Compose(
         [
-            T.Resize(resize_size, antialias=True),
+            T.Resize((img_size + 50, img_size + 50), antialias=True),
             T.CenterCrop(img_size),
             T.ToTensor(),
             T.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]),

src/aging_gan/train.py

@@ -34,28 +34,28 @@ def parse_args() -> argparse.Namespace:
     p.add_argument(
         "--gen_lr",
         type=float,
-        default=3e-4,
+        default=1e-4,
         help="Initial learning rate for generators.",
     )
     p.add_argument(
         "--disc_lr",
         type=float,
-        default=2e-4,
+        default=1e-4,
         help="Initial learning rate for discriminators.",
     )
     p.add_argument(
-        "--num_train_epochs", type=int, default=25, help="Number of training epochs."
+        "--num_train_epochs", type=int, default=100, help="Number of training epochs."
     )
     p.add_argument(
         "--train_batch_size",
         type=int,
-        default=16,
+        default=4,
         help="Batch size per device during training.",
     )
     p.add_argument(
         "--eval_batch_size",
         type=int,
-        default=32,
+        default=8,
         help="Batch size per device during evaluation.",
     )
     p.add_argument(
@@ -67,7 +67,7 @@ def parse_args() -> argparse.Namespace:
     p.add_argument(
         "--lambda_cyc_value",
         type=int,
-        default=7,
+        default=10,
         help="Weight for cyclical loss",
     )
     p.add_argument(
@@ -164,7 +164,7 @@ def initialize_optimizers(cfg, G, F, DX, DY):
 
 
 def initialize_loss_functions(
-    lambda_adv_value: int = 2, lambda_cyc_value: int = 10, lambda_id_value: int = 5
+    lambda_adv_value: int = 2, lambda_cyc_value: int = 10, lambda_id_value: int = 7
 ):
     mse = nn.MSELoss()
     l1 = nn.L1Loss()
@@ -643,22 +643,23 @@ def main() -> None:
                 "best",
             )
         # save the latest checkpoint
-        save_checkpoint(
-            epoch,
-            G,
-            F,
-            DX,
-            DY,
-            opt_G,
-            opt_F,  # generator optimizers
-            opt_DX,
-            opt_DY,  # discriminator optimizers
-            sched_G,
-            sched_F,
-            sched_DX,
-            sched_DY,  # schedulers
-            "latest",
-        )
+        if epoch % 5 == 0:
+            save_checkpoint(
+                epoch,
+                G,
+                F,
+                DX,
+                DY,
+                opt_G,
+                opt_F,  # generator optimizers
+                opt_DX,
+                opt_DY,  # discriminator optimizers
+                sched_G,
+                sched_F,
+                sched_DX,
+                sched_DY,  # schedulers
+                "current",
+            )
 
     # ---------- Test ----------
     if cfg.do_test:

src/aging_gan/utils.py

@@ -64,7 +64,7 @@ def save_checkpoint(
     sched_DY,  # schedulers
     kind: str = "best",
 ):
-    """Overwrite the single best‐ever checkpoint."""
+    """Overwrite the single best-ever checkpoint."""
     ckpt_dir = Path(__file__).resolve().parents[2] / "outputs/checkpoints"
     os.makedirs(ckpt_dir, exist_ok=True)
 
@@ -88,15 +88,9 @@ def save_checkpoint(
         filename = os.path.join(ckpt_dir, "best.pth")
         torch.save(state, filename)
         logger.info(f"Saved best checkpoint: {filename}")
-    elif kind == "latest":
+    elif kind == "current":
         new_latest = ckpt_dir / f"epoch_{epoch:04d}.pth"
         torch.save(state, new_latest)
-
-        # remove previous epoch_*.pth checkpoints
-        for f in ckpt_dir.glob("epoch_*.pth"):
-            if f != new_latest:
-                f.unlink(missing_ok=True)
-
         logger.info(f"Saved latest checkpoint: {new_latest}")
     else:
         raise ValueError(f"kind must be 'best' or 'latest', got {kind}")