diff options
Diffstat (limited to 'sdil/data.py')
| -rw-r--r-- | sdil/data.py | 78 |
1 files changed, 74 insertions, 4 deletions
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, |