chore: capture initial SDIL project state

1 files changed, 200 insertions, 0 deletions

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