diff options
| -rw-r--r-- | experiments/protocol_smoke.py | 39 | ||||
| -rw-r--r-- | experiments/run.py | 75 | ||||
| -rw-r--r-- | experiments/synthetic_smoke.py | 3 | ||||
| -rw-r--r-- | sdil/data.py | 78 |
4 files changed, 172 insertions, 23 deletions
diff --git a/experiments/protocol_smoke.py b/experiments/protocol_smoke.py new file mode 100644 index 0000000..36da996 --- /dev/null +++ b/experiments/protocol_smoke.py @@ -0,0 +1,39 @@ +"""Fast checks for frozen train/validation/test protocol plumbing.""" +import os +import sys + +import torch + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +from sdil.data import get_dataset_splits + + +def loader_labels(loader): + return torch.cat([labels.cpu() for _, labels in loader]) + + +def main(): + first = get_dataset_splits( + "mnist", batch_size=256, device="cpu", val_examples=1000, split_seed=2027) + second = get_dataset_splits( + "mnist", batch_size=256, device="cpu", val_examples=1000, split_seed=2027) + train, validation, test, n_in, n_out, metadata = first + _, validation_2, _, _, _, metadata_2 = second + + assert (n_in, n_out) == (784, 10) + assert metadata["train_examples"] == 59000 + assert metadata["validation_examples"] == 1000 + assert metadata["test_examples"] == 10000 + assert metadata["validation_index_sha256"] == metadata_2["validation_index_sha256"] + assert metadata["validation_class_counts"] == {str(i): 100 for i in range(10)} + assert torch.equal(loader_labels(validation), loader_labels(validation_2)) + assert sum(y.numel() for _, y in train) == 59000 + assert sum(y.numel() for _, y in validation) == 1000 + assert sum(y.numel() for _, y in test) == 10000 + print("validation split hash:", metadata["validation_index_sha256"]) + print("train/validation/test: 59000/1000/10000; stratification exact") + print("ALL PROTOCOL SMOKE CHECKS PASSED") + + +if __name__ == "__main__": + main() diff --git a/experiments/run.py b/experiments/run.py index cc3e475..c66acca 100644 --- a/experiments/run.py +++ b/experiments/run.py @@ -24,7 +24,7 @@ from sdil.core import SDILNet, SDILConfig, sdil_step, neutral_p_update from sdil.baselines import BPNet, dfa_config, evaluate from sdil.local_baselines import FANet from sdil import probes -from sdil.data import (get_dataset, onehot, make_hierarchical, +from sdil.data import (get_dataset_splits, onehot, make_hierarchical, make_teacher_student, make_tentmap) @@ -95,9 +95,18 @@ def load_task(args, device): teacher between model seeds. """ if args.dataset in REAL_DATASETS: - return get_dataset( + train, validation, test, n_in, n_out, split = get_dataset_splits( args.dataset, batch_size=args.batch_size, device=device, - train_limit=args.train_examples or None) + train_limit=args.train_examples or None, + val_examples=args.val_examples, split_seed=args.split_seed) + if args.eval_split == "validation": + if validation is None: + raise ValueError("--eval_split validation requires --val_examples > 0") + evaluation = validation + else: + evaluation = test + split["evaluation_split"] = args.eval_split + return train, evaluation, n_in, n_out, split common = dict( n_train=args.task_train_examples, n_test=args.task_test_examples, @@ -106,32 +115,46 @@ def load_task(args, device): device=device, ) if args.dataset == "teacher": - return make_teacher_student( + task = make_teacher_student( n_in=args.task_n_in, n_classes=args.task_classes, t_depth=args.teacher_depth, t_width=args.teacher_width, residual=bool(args.teacher_residual), **common) - if args.dataset == "hierarchical": - return make_hierarchical( + elif args.dataset == "hierarchical": + task = make_hierarchical( levels=args.task_levels, n_classes=args.task_classes, **common) - if args.dataset == "tentmap": - return make_tentmap( + elif args.dataset == "tentmap": + task = make_tentmap( levels=args.task_levels, n_in=args.task_n_in, **common) - raise ValueError(f"unknown dataset: {args.dataset}") + else: + raise ValueError(f"unknown dataset: {args.dataset}") + train, test, n_in, n_out = task + split = { + "dataset": args.dataset, + "task_seed": args.task_seed, + "train_examples": args.task_train_examples, + "test_examples": args.task_test_examples, + "evaluation_split": "test", + "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 def train(args): device = args.device torch.manual_seed(args.seed) - train_loader, test_loader, n_in, n_out = load_task(args, device) + train_loader, eval_loader, n_in, n_out, split = load_task(args, device) args.n_in, args.n_out = n_in, n_out net, cfg = build(args, device) # a fixed probe batch for stable alignment tracking - px, py = next(iter(test_loader)) + px, py = next(iter(eval_loader)) px, py = px[:args.probe_bs].to(device), py[:args.probe_bs].to(device) poh = onehot(py, n_out, device=device) - log = {"args": vars(args), "provenance": code_provenance(), "steps": [], "final": {}} + log = {"args": vars(args), "split": split, "provenance": code_provenance(), + "steps": [], "final": {}} step = 0 prev_error = None t0 = time.time() @@ -182,8 +205,9 @@ def train(args): if args.max_steps and step >= args.max_steps: break - acc, tloss = evaluate(net, test_loader) - msg = f"[{args.tag}] epoch {epoch} step {step} loss {loss:.4f} test_acc {acc:.4f}" + acc, tloss = evaluate(net, eval_loader) + metric = "val_acc" if args.eval_split == "validation" else "test_acc" + msg = f"[{args.tag}] epoch {epoch} step {step} loss {loss:.4f} {metric} {acc:.4f}" if args.mode == "fa": al = probes.fa_alignment_report(net, px, py, poh) meancos = sum(al["cos_fa_negg"]) / len(al["cos_fa_negg"]) @@ -193,10 +217,21 @@ def train(args): meancos = sum(al["cos_r_negg"]) / len(al["cos_r_negg"]) msg += f" mean_cos(r,-g) {meancos:+.3f} per-layer {['%.2f'%v for v in al['cos_r_negg']]}" print(msg, flush=True) - log["steps"].append({"epoch_end": epoch, "step": step, "test_acc": acc, "test_loss": tloss}) + record = {"epoch_end": epoch, "step": step, "eval_split": args.eval_split, + "eval_acc": acc, "eval_loss": tloss} + if args.eval_split == "test": + record.update({"test_acc": acc, "test_loss": tloss}) + else: + record.update({"val_acc": acc, "val_loss": tloss}) + log["steps"].append(record) - acc, tloss = evaluate(net, test_loader) - log["final"] = {"test_acc": acc, "test_loss": tloss, "wall_s": time.time() - t0} + acc, tloss = evaluate(net, eval_loader) + log["final"] = {"eval_split": args.eval_split, "eval_acc": acc, "eval_loss": tloss, + "wall_s": time.time() - t0} + if args.eval_split == "test": + log["final"].update({"test_acc": acc, "test_loss": tloss}) + else: + log["final"].update({"val_acc": acc, "val_loss": tloss}) if args.mode == "fa": log["final"].update(probes.fa_alignment_report(net, px, py, poh)) elif args.mode != "bp": @@ -212,7 +247,7 @@ def train(args): outpath = os.path.join(args.outdir, f"{args.tag}.json") with open(outpath, "w") as f: json.dump(log, f) - print(f"[{args.tag}] DONE test_acc={acc:.4f} -> {outpath}", flush=True) + print(f"[{args.tag}] DONE {args.eval_split}_acc={acc:.4f} -> {outpath}", flush=True) return log @@ -229,6 +264,10 @@ def get_args(): p.add_argument("--batch_size", type=int, default=128) p.add_argument("--train_examples", type=int, default=0, help="0 uses the full training split") + p.add_argument("--val_examples", type=int, default=0, + help="stratified validation examples held out from training") + p.add_argument("--split_seed", type=int, default=2027) + p.add_argument("--eval_split", default="test", choices=["validation", "test"]) p.add_argument("--task_seed", type=int, default=0, help="fixed target/data seed, separate from student --seed") p.add_argument("--task_train_examples", type=int, default=50000) diff --git a/experiments/synthetic_smoke.py b/experiments/synthetic_smoke.py index cd6e2af..d2325a8 100644 --- a/experiments/synthetic_smoke.py +++ b/experiments/synthetic_smoke.py @@ -20,7 +20,7 @@ def check_task(dataset, expected_in, expected_out, extra): args = get_args() finally: sys.argv = old - train, test, n_in, n_out = load_task(args, "cpu") + train, test, n_in, n_out, split = load_task(args, "cpu") assert (n_in, n_out) == (expected_in, expected_out) args.n_in, args.n_out = n_in, n_out x, y = next(iter(train)) @@ -31,6 +31,7 @@ def check_task(dataset, expected_in, expected_out, extra): net, _ = build(args, "cpu") assert net.logits(x).shape == (16, n_out) assert sum(batch_y.numel() for _, batch_y in test) == 32 + assert split["task_seed"] == args.task_seed print(f"{dataset}: input={n_in}, classes={n_out}, fixed task seed={args.task_seed}") 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, |