added tests

Author: codinglabsong

src/aging_gan/__init__.py

@@ -0,0 +1,6 @@
+import sys
+from pathlib import Path
+
+# Add src directory to path for tests
+SRC_ROOT = Path(__file__).resolve().parents[1] / "src"
+sys.path.insert(0, str(SRC_ROOT))

tests/conftest.py

@@ -0,0 +1,46 @@
+from pathlib import Path
+from PIL import Image
+import torchvision.transforms as T
+from aging_gan import data
+
+
+def create_utk_dataset(tmp_path, num_per_split=6):
+    root = Path(tmp_path)
+    ds_root = root / "utkface_aligned_cropped" / "UTKFace"
+    ds_root.mkdir(parents=True)
+    # create young images ages 18..(18+num_per_split-1)
+    for i in range(num_per_split):
+        age = 18 + i
+        img = Image.new("RGB", (32, 32), color=(i, i, i))
+        img.save(ds_root / f"{age}_0_0_202001010000{i}.jpg")
+    for i in range(num_per_split):
+        age = 40 + i
+        img = Image.new("RGB", (32, 32), color=(i, i, i))
+        img.save(ds_root / f"{age}_0_0_202001010100{i}.jpg")
+    return root
+
+
+def test_utkface_len_and_getitem(tmp_path):
+    root = create_utk_dataset(tmp_path)
+    ds = data.UTKFace(str(root))
+    assert len(ds) == 12
+    img, age = ds[0]
+    assert isinstance(age, int)
+    assert isinstance(img, Image.Image)
+
+
+def test_make_unpaired_loader(tmp_path):
+    root = create_utk_dataset(tmp_path)
+    loader = data.make_unpaired_loader(
+        str(root),
+        "train",
+        T.Compose([T.ToTensor()]),
+        batch_size=2,
+        num_workers=1,
+        seed=0,
+        young_max=23,
+        old_min=40,
+    )
+    x, y = next(iter(loader))
+    assert x.shape == y.shape
+    assert x.shape[0] == 2

tests/test_data.py

@@ -0,0 +1,24 @@
+import torch
+from aging_gan import model
+
+
+def test_generator_output_shape():
+    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():
+    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():
+    G, F, DX, DY = model.initialize_models(ngf=8, ndf=8, n_blocks=1)
+    assert isinstance(G, model.Generator)
+    assert isinstance(F, model.Generator)
+    assert isinstance(DX, model.Discriminator)
+    assert isinstance(DY, model.Discriminator)

tests/test_model.py

@@ -0,0 +1,33 @@
+import torch
+from types import SimpleNamespace
+from aging_gan import train, model
+
+
+import sys
+
+
+def test_parse_args_defaults(monkeypatch):
+    monkeypatch.setattr(sys, "argv", ["prog"])
+    args = train.parse_args()
+    assert args.gen_lr == 2e-4
+    assert args.disc_lr == 1e-4
+    assert args.num_train_epochs == 100
+
+
+def test_initialize_loss_functions_defaults():
+    mse, l1, adv, cyc, ident = train.initialize_loss_functions()
+    assert isinstance(mse, torch.nn.MSELoss)
+    assert adv == 2.0
+    assert cyc == 10.0
+    assert ident == 7.0
+
+
+def test_make_schedulers_decay():
+    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]
+    sched_G, _, _, _ = train.make_schedulers(cfg, *opts)
+    sched_G.step()  # epoch 0
+    assert opts[0].param_groups[0]["lr"] == 1.0
+    sched_G.step(3)
+    assert opts[0].param_groups[0]["lr"] < 1.0

tests/test_train.py

@@ -0,0 +1,49 @@
+import random
+import numpy as np
+import torch
+from pathlib import Path
+from aging_gan import utils
+
+
+def test_set_seed_reproducibility():
+    utils.set_seed(123)
+    a = random.random()
+    b = np.random.rand()
+    c = torch.rand(1)
+
+    utils.set_seed(123)
+    assert random.random() == a
+    assert np.random.rand() == b
+    assert torch.allclose(torch.rand(1), c)
+
+
+def test_get_device_cpu():
+    assert utils.get_device().type == "cpu"
+
+
+def test_save_checkpoint(tmp_path):
+    model = torch.nn.Linear(1, 1)
+    opt = torch.optim.SGD(model.parameters(), lr=0.1)
+    sched = torch.optim.lr_scheduler.LambdaLR(opt, lambda _: 1)
+
+    utils.save_checkpoint(
+        1,
+        model,
+        model,
+        model,
+        model,
+        opt,
+        opt,
+        opt,
+        opt,
+        sched,
+        sched,
+        sched,
+        sched,
+        kind="best",
+    )
+    ckpt_file = (
+        Path(__file__).resolve().parents[1] / "outputs" / "checkpoints" / "best.pth"
+    )
+    assert ckpt_file.exists()
+    ckpt_file.unlink()