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
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
|
"""Canonical runners for non-backprop baselines outside the SDIL family.
The methods intentionally do not all share one architecture: Forward-Forward
has goodness layers instead of a classifier readout, while Equilibrium
Propagation is an energy-based recurrent network with symmetric connections.
Forcing either into SDILNet would make the comparison look uniform while
silently changing the algorithm. Every JSON result records the protocol.
"""
import argparse
import json
import os
import subprocess
import sys
import time
import torch
import torch.nn.functional as F
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from sdil.data import get_dataset, onehot
from sdil.local_baselines import FANet, PEPITANet, FFNet, EPNet
METHOD_SOURCES = {
"fa": {
"paper": "https://www.nature.com/articles/ncomms13276",
"protocol": "fixed random sequential feedback matrices",
},
"pepita": {
"paper": "https://proceedings.mlr.press/v162/dellaferrera22a.html",
"code": "https://github.com/GiorgiaD/PEPITA",
"protocol": "ERIN two-forward-pass rule; He-uniform; ReLU; F scale 0.05",
},
"ff": {
"paper": "https://www.cs.toronto.edu/~hinton/absps/FFXfinal.pdf",
"reference": "https://github.com/mpezeshki/pytorch_forward_forward",
"protocol": "greedy layerwise goodness; input length normalization; Adam",
},
"ep": {
"paper": "https://www.frontiersin.org/journals/computational-neuroscience/articles/10.3389/fncom.2017.00024/full",
"code": "https://github.com/bscellier/Towards-a-Biologically-Plausible-Backprop",
"protocol": "two-phase energy-gradient dynamics; random beta sign",
},
}
_MNIST_STATS = {"mnist": (0.1307, 0.3081), "fmnist": (0.2860, 0.3530)}
_CIFAR_MEAN = torch.tensor([0.4914, 0.4822, 0.4465])
_CIFAR_STD = torch.tensor([0.2470, 0.2435, 0.2616])
def code_provenance():
root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
try:
commit = subprocess.run(["git", "rev-parse", "HEAD"], cwd=root, check=True,
capture_output=True, text=True).stdout.strip()
dirty = bool(subprocess.run(
["git", "status", "--porcelain", "--untracked-files=no"], cwd=root,
check=True, capture_output=True, text=True).stdout.strip())
return {"git_commit": commit, "git_dirty": dirty}
except (OSError, subprocess.CalledProcessError):
return {"git_commit": None, "git_dirty": None}
def canonical_input(x, dataset, method, enabled=True):
"""Restore [0,1] pixels where the authors' protocol requires them."""
if not enabled or method not in ("pepita", "ep"):
return x
if dataset in _MNIST_STATS:
mean, std = _MNIST_STATS[dataset]
return (x * std + mean).clamp(0, 1)
if dataset == "cifar10":
mean = _CIFAR_MEAN.to(x.device, x.dtype).view(1, 3, 1)
std = _CIFAR_STD.to(x.device, x.dtype).view(1, 3, 1)
return (x.view(-1, 3, 1024) * std + mean).clamp(0, 1).view_as(x)
return x
@torch.no_grad()
def evaluate_method(net, loader, method, dataset, canonical):
correct = total = 0
loss_sum = 0.0
for x, y in loader:
x = canonical_input(x, dataset, method, canonical)
if method == "ff":
pred = net.predict(x)
elif method == "ep":
state = net._settle(x, beta=0.0, T=net.T_free)
logits = state[-1]
pred = logits.argmax(1)
loss_sum += F.mse_loss(logits, onehot(y, 10), reduction="sum").item()
else:
logits = net.logits(x)
pred = logits.argmax(1)
loss_sum += F.cross_entropy(logits, y, reduction="sum").item()
correct += (pred == y).sum().item()
total += y.shape[0]
return correct / total, loss_sum / total if method != "ff" else float("nan")
def ep_learning_rates(depth, base_eta=None):
"""Layerwise rates from the authors' one/two/three-hidden-layer runs."""
canonical = {
1: [0.1, 0.05],
2: [0.4, 0.1, 0.01],
3: [0.128, 0.032, 0.008, 0.002],
}
if base_eta is not None:
return [base_eta / (4 ** l) for l in range(depth + 1)]
if depth in canonical:
return canonical[depth]
# The original implementation already required rapidly shrinking rates as
# depth grew. Continue its depth-3 geometric schedule for an explicit,
# logged extrapolation rather than pretending there is a canonical one.
return [0.128 / (4 ** l) for l in range(depth + 1)]
def build(args, n_in, device):
sizes = [n_in] + [args.width] * args.depth + [10]
if args.method == "fa":
return FANet(sizes, act=args.act, device=device, seed=args.seed,
residual=bool(args.residual), b_scale=args.feedback_scale)
if args.method == "pepita":
return PEPITANet(sizes, act="relu", device=device, seed=args.seed,
f_scale=args.pepita_f_scale, keep_prob=args.keep_prob)
if args.method == "ff":
# FF has only goodness layers; a 10-unit SDIL-style output is not part
# of the original supervised algorithm.
return FFNet([n_in] + [args.width] * args.depth, device=device,
seed=args.seed, threshold=args.ff_threshold)
if args.method == "ep":
return EPNet(sizes, device=device, seed=args.seed, beta=args.ep_beta,
dt=args.ep_dt, T_free=args.ep_free_steps,
T_nudge=args.ep_nudge_steps, random_beta_sign=True)
raise ValueError(args.method)
def train(args):
torch.manual_seed(args.seed)
device = args.device
train_loader, test_loader, n_in, n_out = get_dataset(
args.dataset, args.batch_size, device=device)
net = build(args, n_in, device)
canonical = bool(args.canonical_preprocess)
log = {
"args": vars(args),
"method_source": METHOD_SOURCES[args.method],
"provenance": code_provenance(),
"steps": [],
"final": {},
}
t0 = time.time()
if args.method == "ff":
eta = 0.03 if args.eta is None else args.eta
for layer in range(args.depth):
for epoch in range(args.epochs):
batches = 0
for x, y in train_loader:
loss = net.train_layer(layer, x, y, eta)
batches += 1
if args.max_batches and batches >= args.max_batches:
break
print(f"[{args.tag}] layer {layer} epoch {epoch} local_loss {loss:.4f}",
flush=True)
acc, test_loss = evaluate_method(net, test_loader, args.method,
args.dataset, canonical)
log["steps"].append({"layer_end": layer, "test_acc": acc,
"local_loss": loss})
print(f"[{args.tag}] layer {layer} test_acc {acc:.4f}", flush=True)
else:
eta = args.eta
if eta is None:
eta = 0.05 if args.method == "fa" else (0.1 if args.method == "pepita" else None)
ep_etas = ep_learning_rates(args.depth, eta) if args.method == "ep" else None
for epoch in range(args.epochs):
if args.method == "pepita" and epoch in (60, 90):
eta *= 0.1
batches = 0
for x, y in train_loader:
x = canonical_input(x, args.dataset, args.method, canonical)
yoh = onehot(y, n_out, device=device)
if args.method == "fa":
loss = net.fa_step(x, y, yoh, eta, args.momentum)
elif args.method == "pepita":
loss = net.pepita_step(x, y, yoh, eta, args.momentum)
else:
loss = net.train_step(x, y, yoh, ep_etas)
batches += 1
if args.max_batches and batches >= args.max_batches:
break
acc, test_loss = evaluate_method(net, test_loader, args.method,
args.dataset, canonical)
log["steps"].append({"epoch_end": epoch, "test_acc": acc,
"test_loss": test_loss, "train_loss": loss})
print(f"[{args.tag}] epoch {epoch} loss {loss:.4f} test_acc {acc:.4f}",
flush=True)
acc, test_loss = evaluate_method(net, test_loader, args.method,
args.dataset, canonical)
log["final"] = {"test_acc": acc, "test_loss": test_loss,
"wall_s": time.time() - t0}
os.makedirs(args.outdir, exist_ok=True)
path = os.path.join(args.outdir, f"{args.tag}.json")
with open(path, "w") as f:
json.dump(log, f)
print(f"[{args.tag}] DONE test_acc={acc:.4f} -> {path}", flush=True)
return log
def get_args():
p = argparse.ArgumentParser()
p.add_argument("--method", required=True, choices=["fa", "pepita", "ff", "ep"])
p.add_argument("--dataset", default="mnist", choices=["mnist", "fmnist", "cifar10"])
p.add_argument("--depth", type=int, default=2)
p.add_argument("--width", type=int, default=500)
p.add_argument("--epochs", type=int, default=10,
help="FF: epochs per greedy layer; other methods: global epochs")
p.add_argument("--batch_size", type=int, default=64)
p.add_argument("--max_batches", type=int, default=0)
p.add_argument("--eta", type=float, default=None)
p.add_argument("--momentum", type=float, default=0.9)
p.add_argument("--act", default="tanh", choices=["tanh", "gelu", "silu", "relu"])
p.add_argument("--residual", type=int, default=0)
p.add_argument("--feedback_scale", type=float, default=1.0)
p.add_argument("--pepita_f_scale", type=float, default=0.05)
p.add_argument("--keep_prob", type=float, default=0.9)
p.add_argument("--ff_threshold", type=float, default=2.0)
p.add_argument("--ep_beta", type=float, default=0.5)
p.add_argument("--ep_dt", type=float, default=0.5)
p.add_argument("--ep_free_steps", type=int, default=20)
p.add_argument("--ep_nudge_steps", type=int, default=4)
p.add_argument("--canonical_preprocess", type=int, default=1)
p.add_argument("--seed", type=int, default=0)
p.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
p.add_argument("--outdir", default="results")
p.add_argument("--tag", default="baseline")
return p.parse_args()
if __name__ == "__main__":
train(get_args())
|
