From 215acff518eea2d6d8a5bd7c90c398f49cc3b41b Mon Sep 17 00:00:00 2001 From: YurenHao0426 Date: Tue, 21 Jul 2026 06:38:55 -0500 Subject: chore: capture initial SDIL project state --- sdil/data.py | 200 +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 200 insertions(+) create mode 100644 sdil/data.py (limited to 'sdil/data.py') diff --git a/sdil/data.py b/sdil/data.py new file mode 100644 index 0000000..f5b4928 --- /dev/null +++ b/sdil/data.py @@ -0,0 +1,200 @@ +"""Data loaders. Parse the raw IDX files directly (no torchvision dependency, +so it runs on the Pascal env on the 1080 farm). Datasets already live on the +shared NFS home; the box is offline for dataset mirrors.""" + +import pickle +import struct +import numpy as np +import torch + +DATA_DIR = "/home/yurenh2/sdrn/data" + +# (subdir, mean, std) matching the standard torchvision normalisation used by +# the predecessor project, so numbers are comparable. +_STATS = { + "mnist": ("MNIST", 0.1307, 0.3081), + "fmnist": ("FashionMNIST", 0.2860, 0.3530), +} + +_CIFAR_MEAN = np.array([0.4914, 0.4822, 0.4465], dtype=np.float32) +_CIFAR_STD = np.array([0.2470, 0.2435, 0.2616], dtype=np.float32) + + +def _read_images(path): + with open(path, "rb") as f: + magic, n, r, c = struct.unpack(">IIII", f.read(16)) + assert magic == 2051, f"bad image magic {magic}" + buf = f.read(n * r * c) + return np.frombuffer(buf, dtype=np.uint8).reshape(n, r * c).astype(np.float32) + + +def _read_labels(path): + with open(path, "rb") as f: + magic, n = struct.unpack(">II", f.read(8)) + assert magic == 2049, f"bad label magic {magic}" + buf = f.read(n) + return np.frombuffer(buf, dtype=np.uint8).astype(np.int64) + + +def _load_split(subdir, train, mean, std, data_dir): + raw = f"{data_dir}/{subdir}/raw" + pre = "train" if train else "t10k" + x = _read_images(f"{raw}/{pre}-images-idx3-ubyte") + y = _read_labels(f"{raw}/{pre}-labels-idx1-ubyte") + x = (x / 255.0 - mean) / std # normalise + flatten (already 784) + return torch.from_numpy(x), torch.from_numpy(y) + + +class _FastLoader: + """Minimal shuffled minibatch iterator over in-memory (GPU) tensors.""" + def __init__(self, x, y, batch_size, shuffle, seed=0): + self.x, self.y = x, y + self.bs = batch_size + self.shuffle = shuffle + self.n = x.shape[0] + self.g = torch.Generator(device="cpu").manual_seed(seed) + + def __iter__(self): + idx = torch.randperm(self.n, generator=self.g) if self.shuffle else torch.arange(self.n) + for i in range(0, self.n, self.bs): + j = idx[i:i + self.bs] + yield self.x[j], self.y[j] + + def __len__(self): + return (self.n + self.bs - 1) // self.bs + + +def _load_cifar(data_dir, n_classes=10): + root = f"{data_dir}/cifar-10-batches-py" + def unpickle(fn): + with open(fn, "rb") as f: + return pickle.load(f, encoding="bytes") + xs, ys = [], [] + for i in range(1, 6): + d = unpickle(f"{root}/data_batch_{i}") + xs.append(d[b"data"]); ys.append(np.array(d[b"labels"])) + xtr = np.concatenate(xs).astype(np.float32); ytr = np.concatenate(ys).astype(np.int64) + dte = unpickle(f"{root}/test_batch") + xte = np.array(dte[b"data"], dtype=np.float32); yte = np.array(dte[b"labels"], dtype=np.int64) + + def norm(x): # x: (N, 3072) as R(1024)G(1024)B(1024) + x = x.reshape(-1, 3, 1024) / 255.0 + x = (x - _CIFAR_MEAN[None, :, None]) / _CIFAR_STD[None, :, None] + return x.reshape(-1, 3072).astype(np.float32) + return (torch.from_numpy(norm(xtr)), torch.from_numpy(ytr), + 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): + if name == "cifar10": + xtr, ytr, xte, yte = _load_cifar(data_dir) + n_in = 3072 + else: + subdir, mean, std = _STATS[name] + xtr, ytr = _load_split(subdir, True, mean, std, data_dir) + xte, yte = _load_split(subdir, False, mean, std, data_dir) + n_in = 784 + xtr, ytr = xtr.to(device), ytr.to(device) + xte, yte = xte.to(device), yte.to(device) + return (_FastLoader(xtr, ytr, batch_size, True), + _FastLoader(xte, yte, 1000, False), n_in, 10) + + +def make_teacher_student(n_in=128, n_classes=10, t_depth=8, t_width=64, + n_train=50000, n_test=10000, seed=0, residual=True, + t_scale=1.5, batch_size=128, device="cpu"): + """Deep teacher-student classification. Labels = argmax of a fixed random + deep (residual) teacher net. Students are WIDTH-LIMITED, so depth is the + resource that matters: a shallow student cannot compose enough nonlinearity + to match a deep teacher, a deep one can. This is the standard construction + for a task where accuracy genuinely GROWS with model depth — unlike + flattened-CIFAR MLPs, which are architecture-bottlenecked and depth-flat.""" + from .core import SDILNet + tsizes = [n_in] + [t_width] * t_depth + [n_classes] + teacher = SDILNet(tsizes, act="tanh", device=device, seed=seed + 999, + w_scale=t_scale, residual=residual) + g = torch.Generator(device="cpu").manual_seed(seed) + xtr = torch.randn(n_train, n_in, generator=g).to(device) + xte = torch.randn(n_test, n_in, generator=g).to(device) + with torch.no_grad(): + ytr = teacher.logits(xtr).argmax(1) + yte = teacher.logits(xte).argmax(1) + return (_FastLoader(xtr, ytr, batch_size, True), + _FastLoader(xte, yte, 1000, False), n_in, n_classes) + + +def make_hierarchical(levels=6, n_classes=10, n_train=50000, n_test=10000, + seed=0, batch_size=128, device="cpu"): + """Hierarchical compositional target with PROVABLE depth-dependence. Input + has 2^levels features; C independent binary trees each fold adjacent pairs + with a fixed random nonlinear combiner v' = tanh(a·l + b·r + c·l·r) (the + l·r product is the depth-hard bit). Label = argmax over the C tree roots. + A depth-d net can only realise ~d tree levels, so accuracy GROWS with depth + up to `levels` — the regime where scaling the model actually helps, and where + a rule with poor deep credit assignment (DFA) cannot follow.""" + n_in = 2 ** levels + g = torch.Generator(device="cpu").manual_seed(seed + 31) + # per-class, per-level, per-node combiner params + combiners = [] + width = n_in + for lev in range(levels): + half = width // 2 + combiners.append(torch.randn(n_classes, half, 3, generator=g)) + width = half + + def teacher(x): # x: (N, n_in) + N = x.shape[0] + outs = [] + for cls in range(n_classes): + v = x + w = n_in + for lev in range(levels): + half = w // 2 + l = v[:, 0:2 * half:2] + r = v[:, 1:2 * half:2] + p = combiners[lev][cls].to(x.device) # (half, 3) + v = torch.tanh(l * p[:, 0] + r * p[:, 1] + (l * r) * p[:, 2]) + w = half + outs.append(v[:, 0]) + return torch.stack(outs, 1) # (N, C) + + xtr = torch.randn(n_train, n_in, generator=g).to(device) + xte = torch.randn(n_test, n_in, generator=g).to(device) + with torch.no_grad(): + rtr = teacher(xtr) + rte = teacher(xte) + # standardise per-class roots on train stats -> roughly balanced argmax + mu, sd = rtr.mean(0, keepdim=True), rtr.std(0, keepdim=True) + 1e-6 + ytr = ((rtr - mu) / sd).argmax(1) + yte = ((rte - mu) / sd).argmax(1) + return (_FastLoader(xtr, ytr, batch_size, True), + _FastLoader(xte, yte, 1000, False), n_in, n_classes) + + +def make_tentmap(levels=8, n_in=4, n_train=50000, n_test=10000, seed=0, + batch_size=128, device="cpu"): + """Telgarsky depth-separation task. Label = [tent^levels(x0) > 0.5], where + tent(v)=1-|2v-1| self-composed `levels` times produces 2^levels oscillations. + A depth-d ReLU net can realise ~d tent compositions, resolving only ~2^d of + the 2^levels teeth, so TEST ACCURACY PROVABLY RISES WITH DEPTH up to `levels`. + This is the regime where scaling the model genuinely buys accuracy, and where + a rule with poor deep credit assignment (DFA/FA) cannot build the composition + and stalls. Extra input dims are distractors (must be ignored).""" + g = torch.Generator(device="cpu").manual_seed(seed + 7) + xtr = torch.rand(n_train, n_in, generator=g).to(device) + xte = torch.rand(n_test, n_in, generator=g).to(device) + + def label(x): + v = x[:, 0].clone() + for _ in range(levels): + v = 1.0 - (2.0 * v - 1.0).abs() + return (v > 0.5).long() + + return (_FastLoader(xtr, label(xtr), batch_size, True), + _FastLoader(xte, label(xte), 1000, False), n_in, 2) + + +def onehot(y, n_classes, device=None): + oh = torch.zeros(y.shape[0], n_classes, device=y.device if device is None else device) + oh.scatter_(1, y.view(-1, 1), 1.0) + return oh -- cgit v1.2.3