Merge branch 'main' of github.com:codinglabsong/aging-gan

Author: codinglabsong

src/aging_gan/data.py

@@ -12,13 +12,13 @@
 
 class UTKFace(Dataset):
     """
-    Assumes the unzipped UTKFace images live in  <root>/data/UTKFace
+    Assumes the unzipped aligned UTKFace images live in  <root>/data/utkface_aligned_cropped/UTKFace
     File pattern:  {age}_{gender}_{race}_{yyyymmddHHMMSS}.jpg
     """
 
     def __init__(self, root: str, transform: T.Compose | None = None):
         self.root = (
-            Path(root) / "UTKFace"  # "utkface_aligned_cropped" /
+            Path(root) / "utkface_aligned_cropped" / "UTKFace"
         )  # or "UTKFace" for the unaligned and varied original version.
         self.files = sorted(f for f in self.root.glob("*.jpg"))
         if not self.files:
@@ -132,7 +132,7 @@ def prepare_dataset(
     # randomness
     train_transform = T.Compose(
         [
-            T.ToPILImage(),
+            # T.ToPILImage(),
             T.RandomHorizontalFlip(),
             T.Resize((img_size + 50, img_size + 50), antialias=True),
             T.RandomCrop(img_size),

src/aging_gan/train.py

@@ -94,7 +94,7 @@ def parse_args() -> argparse.Namespace:
     p.add_argument(
         "--steps_for_logging_metrics",
         type=int,
-        default=1,
+        default=50,
         help="Print training metrics after certain batch steps.",
     )
     p.add_argument(

src/aging_gan/utils.py

@@ -103,9 +103,18 @@ def generate_and_save_samples(
     device: torch.device,
     num_samples: int = 8,
 ):
-    # grab one batch
-    inputs, _ = next(iter(val_loader))
-    inputs = inputs.to(device)[:num_samples]
+    # grab batches until num_samples   
+    collected = []
+    for imgs, _ in val_loader:
+        collected.append(imgs)
+        if sum(b.size(0) for b in collected) >= num_samples:
+            break
+
+    if not collected:
+        raise ValueError("Validation loader is empty.")
+    
+    inputs = torch.cat(collected, dim=0)[:num_samples].to(device)
+    
     with torch.no_grad():
         outputs = generator(inputs)