1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
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
|
