"""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