added annotations for tests

Author: codinglabsong

README.md

@@ -1,22 +1,20 @@
 different lr for gen and disc
 
-
 # Aging GAN
 
-Aging GAN is a research project exploring face aging with a CycleGAN-style architecture. The code trains two U-Net generators and two PatchGAN discriminators on the CelebA dataset, preprocessing into **Young** and **Old** subsets. The generators learn to translate between these domains, effectively "aging" or "de-aging" a face image.
+Aging GAN is a research project exploring facial age transformation with a CycleGAN‑style approach. The model trains two ResNet‑style "encoder–residual–decoder" generators and two PatchGAN discriminators on the UTKFace dataset, split into **Young** and **Old** subsets. The generators learn to translate between these domains, effectively "aging" or "de-aging" a face image.
 
-This repository contains training scripts, minimal utilities, and example notebooks.
+This repository contains training scripts, helper utilities, and inference scripts.
 
 ## Features
 
-- **Unpaired Training Data** – automatically split the CelebA dataset into Young vs. Old and create an unpaired `DataLoader`.
-- **CycleGAN Architecture** – U-Net generators with ResNet encoders and PatchGAN discriminators.
-- **Training Utilities** – gradient clipping, learning-rate scheduling, mixed-precision via `accelerate`, and optional encoder freezing.
-- **Evaluation** – FID metric computation on the validation set.
+- **Unpaired Training Data** – splits the UTKFace dataset into *Young* (18‑28) and *Old* (40+) subsets and builds an unpaired `DataLoader`.
+- **CycleGAN Architecture** – residual U‑Net generators and PatchGAN discriminators.
+- **Training Utilities** – gradient clipping, separate generator/discriminator learning rates with linear decay, mixed precision via `accelerate`, and optional S3 checkpoint archiving.
+- **Evaluation** – FID metric computation on the validation and test sets.
 - **Weights & Biases Logging** – track losses and metrics during training.
-- **Scriptable Workflows** – run training from the command line with `scripts/run_train.sh`.
-
-*Placeholders:* inference helpers, web demo, and quantitative results will be added later.
+- **Scriptable Workflows** – shell scripts for training and inference.
+- **Sample Generation** – saves example outputs after each epoch.
 
 ## Installation
 
@@ -38,7 +36,9 @@ pip install -e .
 ```
 
 ## Data
-The `prepare_dataset` function downloads CelebA automatically and creates train, validation, and test splits. Images are center‑cropped and resized to 256×256. Each split is divided into *Young* and *Old* subsets for unpaired training.
+Place the aligned UTKFace images under `data/utkface_aligned_cropped/UTKFace`.
+The `prepare_dataset` function builds deterministic train/val/test splits and applies random flipping, cropping and rotation for training. 
+Each split is divided into *Young* and *Old* subsets for unpaired training.
 
 ## Training
 Run training with default hyper‑parameters:
@@ -54,7 +54,15 @@ python -m aging_gan.train --help
 ```
 
 ## Inference
-The `aging_gan.inference` module is currently a stub. Once implemented, you will be able to generate aged faces from the command line using `scripts/run_inference.sh`.
+Generate aged faces using the command-line helper:
+
+```bash
+bash scripts/run_inference.sh --input myface.jpg --direction young2old
+```
+The script loads `outputs/checkpoints/best.pth` by default and saves the result beside the input.
+
+## AWS Utilities
+When running on EC2, pass `--archive_and_terminate_ec2` to automatically sync `outputs/` to S3 and terminate the instance after training.
 
 ## Results
 *Results will be added here once experiments are complete.*
@@ -65,7 +73,7 @@ The `aging_gan.inference` module is currently a stub. Once implemented, you will
 ## Repository Structure
 
 - `src/aging_gan/` – core modules (`train.py`, `model.py`, etc.)
-- `scripts/` – helper shell scripts for training and (placeholder) inference
+- `scripts/` – helper scripts for training and inference
 - `notebooks/` – exploratory notebooks
 
 ## Requirements

tests/test_data.py

@@ -5,6 +5,7 @@
 
 
 def create_utk_dataset(tmp_path, num_per_split=6):
+    """Create a tiny UTKFace-style directory for testing."""
     root = Path(tmp_path)
     ds_root = root / "utkface_aligned_cropped" / "UTKFace"
     ds_root.mkdir(parents=True)
@@ -21,6 +22,7 @@ def create_utk_dataset(tmp_path, num_per_split=6):
 
 
 def test_utkface_len_and_getitem(tmp_path):
+    """UTKFace yields image/age pairs and correct length."""
     root = create_utk_dataset(tmp_path)
     ds = data.UTKFace(str(root))
     assert len(ds) == 12
@@ -30,6 +32,7 @@ def test_utkface_len_and_getitem(tmp_path):
 
 
 def test_make_unpaired_loader(tmp_path):
+    """Loader returns equal-sized batches of young and old images."""
     root = create_utk_dataset(tmp_path)
     loader = data.make_unpaired_loader(
         str(root),

tests/test_model.py

@@ -3,20 +3,23 @@
 
 
 def test_generator_output_shape():
+    """Generator preserves input image dimensions."""
     G = model.Generator(ngf=8, n_residual_blocks=1)
     x = torch.randn(2, 3, 64, 64)
     y = G(x)
     assert y.shape == x.shape
 
 
 def test_discriminator_output_shape():
+    """Discriminator outputs a single logit per image."""
     D = model.Discriminator(ndf=8)
     x = torch.randn(2, 3, 64, 64)
     out = D(x)
     assert out.shape == (2, 1)
 
 
 def test_initialize_models_types():
+    """Model initializer returns correct component classes."""
     G, F, DX, DY = model.initialize_models(ngf=8, ndf=8, n_blocks=1)
     assert isinstance(G, model.Generator)
     assert isinstance(F, model.Generator)

tests/test_train.py

@@ -7,6 +7,7 @@
 
 
 def test_parse_args_defaults(monkeypatch):
+    """CLI parser returns expected default arguments."""
     monkeypatch.setattr(sys, "argv", ["prog"])
     args = train.parse_args()
     assert args.gen_lr == 2e-4
@@ -15,6 +16,7 @@ def test_parse_args_defaults(monkeypatch):
 
 
 def test_initialize_loss_functions_defaults():
+    """Loss initializer provides default weights and criteria."""
     mse, l1, adv, cyc, ident = train.initialize_loss_functions()
     assert isinstance(mse, torch.nn.MSELoss)
     assert adv == 2.0
@@ -23,6 +25,7 @@ def test_initialize_loss_functions_defaults():
 
 
 def test_make_schedulers_decay():
+    """Learning rate scheduler should decrease learning rate."""
     cfg = SimpleNamespace(num_train_epochs=4)
     models = model.initialize_models(ngf=8, ndf=8, n_blocks=1)
     opts = [torch.optim.SGD(m.parameters(), lr=1.0) for m in models]

tests/test_utils.py

@@ -6,6 +6,7 @@
 
 
 def test_set_seed_reproducibility():
+    """Seeding should make RNGs deterministic."""
     utils.set_seed(123)
     a = random.random()
     b = np.random.rand()
@@ -18,10 +19,12 @@ def test_set_seed_reproducibility():
 
 
 def test_get_device_cpu():
+    """Utility returns CPU when CUDA is unavailable."""
     assert utils.get_device().type == "cpu"
 
 
 def test_save_checkpoint(tmp_path):
+    """Checkpoint file is created on disk."""
     model = torch.nn.Linear(1, 1)
     opt = torch.optim.SGD(model.parameters(), lr=0.1)
     sched = torch.optim.lr_scheduler.LambdaLR(opt, lambda _: 1)