experiments: add synthetic validation isolation

3 files changed, 65 insertions, 6 deletions

diff --git a/experiments/run.py b/experiments/run.py
index c66acca..8c59a50 100644
--- a/experiments/run.py
+++ b/experiments/run.py
@@ -25,7 +25,8 @@ from sdil.baselines import BPNet, dfa_config, evaluate
 from sdil.local_baselines import FANet
 from sdil import probes
 from sdil.data import (get_dataset_splits, onehot, make_hierarchical,
-                       make_teacher_student, make_tentmap)
+                       make_teacher_student, make_tentmap,
+                       split_training_loader)
 
 
 REAL_DATASETS = ("mnist", "fmnist", "cifar10")
@@ -128,17 +129,27 @@ def load_task(args, device):
     else:
         raise ValueError(f"unknown dataset: {args.dataset}")
     train, test, n_in, n_out = task
+    validation = None
+    split_details = {}
+    if args.val_examples:
+        train, validation, split_details = split_training_loader(
+            train, args.val_examples, args.split_seed, args.batch_size)
+    if args.eval_split == "validation":
+        if validation is None:
+            raise ValueError("--eval_split validation requires --val_examples > 0")
+        evaluation = validation
+    else:
+        evaluation = test
     split = {
         "dataset": args.dataset,
         "task_seed": args.task_seed,
-        "train_examples": args.task_train_examples,
+        "train_examples": len(train.x),
         "test_examples": args.task_test_examples,
-        "evaluation_split": "test",
+        "evaluation_split": args.eval_split,
         "synthetic_generator": True,
     }
-    if args.eval_split != "test" or args.val_examples:
-        raise ValueError("synthetic validation splits are not implemented yet")
-    return train, test, n_in, n_out, split
+    split.update(split_details)
+    return train, evaluation, n_in, n_out, split
 
 
 def train(args):
diff --git a/experiments/synthetic_smoke.py b/experiments/synthetic_smoke.py
index d2325a8..2b66892 100644
--- a/experiments/synthetic_smoke.py
+++ b/experiments/synthetic_smoke.py
@@ -35,6 +35,30 @@ def check_task(dataset, expected_in, expected_out, extra):
     print(f"{dataset}: input={n_in}, classes={n_out}, fixed task seed={args.task_seed}")
 
 
+def check_synthetic_validation():
+    argv = [
+        "run.py", "--dataset", "teacher", "--device", "cpu",
+        "--task_train_examples", "80", "--task_test_examples", "32",
+        "--val_examples", "20", "--eval_split", "validation",
+        "--batch_size", "16", "--depth", "2", "--width", "12",
+        "--task_n_in", "20", "--task_classes", "4",
+        "--teacher_depth", "3", "--teacher_width", "10",
+    ]
+    old = sys.argv
+    try:
+        sys.argv = argv
+        args = get_args()
+    finally:
+        sys.argv = old
+    train, validation, _, _, split = load_task(args, "cpu")
+    assert sum(y.numel() for _, y in train) == 60
+    assert sum(y.numel() for _, y in validation) == 20
+    assert split["evaluation_split"] == "validation"
+    assert split["validation_index_sha256"]
+    assert split["split_from_training_only"] is True
+    print("synthetic validation: 60/20; test generator untouched")
+
+
 if __name__ == "__main__":
     torch.manual_seed(0)
     check_task("teacher", 20, 4,
@@ -44,4 +68,5 @@ if __name__ == "__main__":
                ["--task_levels", "4", "--task_classes", "4"])
     check_task("tentmap", 5, 2,
                ["--task_levels", "4", "--task_n_in", "5"])
+    check_synthetic_validation()
     print("ALL SYNTHETIC TASK CHECKS PASSED")
diff --git a/sdil/data.py b/sdil/data.py
index 8dac4df..ad51f59 100644
--- a/sdil/data.py
+++ b/sdil/data.py
@@ -114,6 +114,29 @@ def _stratified_validation_indices(y, n_val, seed):
     return train_idx, val_idx
 
 
+def split_training_loader(loader, val_examples, split_seed, train_batch_size):
+    """Split an in-memory synthetic training loader without touching its test set."""
+    train_idx, val_idx = _stratified_validation_indices(
+        loader.y.detach().cpu(), val_examples, split_seed)
+    train_device_idx = train_idx.to(loader.x.device)
+    val_device_idx = val_idx.to(loader.x.device)
+    train = _FastLoader(loader.x[train_device_idx], loader.y[train_device_idx],
+                        train_batch_size, True)
+    validation = _FastLoader(loader.x[val_device_idx], loader.y[val_device_idx],
+                             1000, False)
+    digest = hashlib.sha256(val_idx.numpy().tobytes()).hexdigest()
+    counts = {str(int(cls)): int((loader.y.detach().cpu()[val_idx] == cls).sum())
+              for cls in torch.unique(loader.y.detach().cpu()[val_idx], sorted=True)}
+    return train, validation, {
+        "split_seed": split_seed,
+        "validation_examples": int(val_examples),
+        "validation_index_sha256": digest,
+        "validation_class_counts": counts,
+        "train_examples": int(train_idx.numel()),
+        "split_from_training_only": True,
+    }
+
+
 def get_dataset_splits(name="mnist", batch_size=128, data_dir=DATA_DIR, device="cpu",
                        shuffle_train=True, train_limit=None, val_examples=0, split_seed=2027):
     """Load real data with an optional training-only validation split.