From 98dddcfabb39c13f652b3aa44f0926ee04428c06 Mon Sep 17 00:00:00 2001
From: YurenHao0426 <Blackhao0426@gmail.com>
Date: Wed, 22 Jul 2026 01:38:39 -0500
Subject: experiments: isolate frozen validation protocols

---
 sdil/data.py | 78 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++----
 1 file changed, 74 insertions(+), 4 deletions(-)

(limited to 'sdil/data.py')

diff --git a/sdil/data.py b/sdil/data.py
index 8253d84..8dac4df 100644
--- a/sdil/data.py
+++ b/sdil/data.py
@@ -4,6 +4,7 @@ shared NFS home; the box is offline for dataset mirrors."""
 
 import pickle
 import struct
+import hashlib
 import numpy as np
 import torch
 
@@ -85,8 +86,42 @@ def _load_cifar(data_dir, n_classes=10):
             torch.from_numpy(norm(xte)), torch.from_numpy(yte))
 
 
-def get_dataset(name="mnist", batch_size=128, data_dir=DATA_DIR, device="cpu", preload=True,
-                shuffle_train=True, train_limit=None):
+def _stratified_validation_indices(y, n_val, seed):
+    """Return deterministic disjoint train/validation indices.
+
+    Validation examples are allocated as evenly as possible across observed
+    classes. The returned validation indices are sorted before hashing so the
+    split identity is independent of loader iteration order.
+    """
+    if n_val <= 0 or n_val >= y.numel():
+        raise ValueError(f"n_val must be in [1, {y.numel() - 1}], got {n_val}")
+    classes = torch.unique(y, sorted=True)
+    base, remainder = divmod(n_val, classes.numel())
+    generator = torch.Generator(device="cpu").manual_seed(seed)
+    val_parts = []
+    for position, cls in enumerate(classes):
+        candidates = torch.nonzero(y == cls, as_tuple=False).flatten()
+        take = base + int(position < remainder)
+        if take > candidates.numel():
+            raise ValueError(
+                f"class {int(cls)} has only {candidates.numel()} examples, requested {take}")
+        order = torch.randperm(candidates.numel(), generator=generator)
+        val_parts.append(candidates[order[:take]])
+    val_idx = torch.cat(val_parts).sort().values
+    is_val = torch.zeros(y.numel(), dtype=torch.bool)
+    is_val[val_idx] = True
+    train_idx = torch.nonzero(~is_val, as_tuple=False).flatten()
+    return train_idx, val_idx
+
+
+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.
+
+    Returns ``(train, validation_or_none, test, n_in, n_out, metadata)``. Test
+    examples never enter split construction. ``get_dataset`` below preserves
+    the historical four-item API for all pre-freeze experiments.
+    """
     if name == "cifar10":
         xtr, ytr, xte, yte = _load_cifar(data_dir)
         n_in = 3072
@@ -95,14 +130,49 @@ def get_dataset(name="mnist", batch_size=128, data_dir=DATA_DIR, device="cpu", p
         xtr, ytr = _load_split(subdir, True, mean, std, data_dir)
         xte, yte = _load_split(subdir, False, mean, std, data_dir)
         n_in = 784
+
+    val_loader = None
+    if val_examples:
+        train_idx, val_idx = _stratified_validation_indices(ytr, val_examples, split_seed)
+        xval, yval = xtr[val_idx], ytr[val_idx]
+        xtr, ytr = xtr[train_idx], ytr[train_idx]
+        digest = hashlib.sha256(val_idx.numpy().tobytes()).hexdigest()
+        val_counts = {str(int(cls)): int((yval == cls).sum())
+                      for cls in torch.unique(yval, sorted=True)}
+    else:
+        val_idx = None
+        xval = yval = None
+        digest = None
+        val_counts = {}
     if train_limit is not None:
         if train_limit <= 0 or train_limit > xtr.shape[0]:
             raise ValueError(f"train_limit must be in [1, {xtr.shape[0]}], got {train_limit}")
         xtr, ytr = xtr[:train_limit], ytr[:train_limit]
     xtr, ytr = xtr.to(device), ytr.to(device)
     xte, yte = xte.to(device), yte.to(device)
-    return (_FastLoader(xtr, ytr, batch_size, shuffle_train),
-            _FastLoader(xte, yte, 1000, False), n_in, 10)
+    if xval is not None:
+        xval, yval = xval.to(device), yval.to(device)
+        val_loader = _FastLoader(xval, yval, 1000, False)
+    metadata = {
+        "dataset": name,
+        "split_seed": split_seed if val_examples else None,
+        "validation_examples": int(val_examples),
+        "validation_index_sha256": digest,
+        "validation_class_counts": val_counts,
+        "train_examples": int(xtr.shape[0]),
+        "test_examples": int(xte.shape[0]),
+        "split_from_training_only": True,
+    }
+    return (_FastLoader(xtr, ytr, batch_size, shuffle_train), val_loader,
+            _FastLoader(xte, yte, 1000, False), n_in, 10, metadata)
+
+
+def get_dataset(name="mnist", batch_size=128, data_dir=DATA_DIR, device="cpu", preload=True,
+                shuffle_train=True, train_limit=None):
+    train, _, test, n_in, n_out, _ = get_dataset_splits(
+        name=name, batch_size=batch_size, data_dir=data_dir, device=device,
+        shuffle_train=shuffle_train, train_limit=train_limit)
+    return train, test, n_in, n_out
 
 
 def make_teacher_student(n_in=128, n_classes=10, t_depth=8, t_width=64,
-- 
cgit v1.2.3