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
|
"""Baselines sharing SDILNet's exact architecture and initialisation, so
comparisons are apples-to-apples.
- BP: exact backprop (autograd) -- the alignment/performance upper bound.
- DFA is NOT a separate class: it is exactly SDIL with a fixed-random apical
vectorizer, no residual, no predictor, no perturbation. See dfa_config().
"""
import torch
import torch.nn.functional as F
from .core import SDILNet, SDILConfig
class BPNet(SDILNet):
"""Same tanh MLP, trained by exact backprop via autograd."""
def bp_step(self, x, y, eta, momentum=0.0):
for W in self.W:
W.requires_grad_(True)
for b in self.b:
b.requires_grad_(True)
logits = self.logits(x)
loss = F.cross_entropy(logits, y)
grads = torch.autograd.grad(loss, self.W + self.b)
nW = len(self.W)
with torch.no_grad():
for i in range(nW):
if momentum:
self.mW[i].mul_(momentum).add_(grads[i])
self.mb[i].mul_(momentum).add_(grads[nW + i])
self.W[i] -= eta * self.mW[i]
self.b[i] -= eta * self.mb[i]
else:
self.W[i] -= eta * grads[i]
self.b[i] -= eta * grads[nW + i]
for W in self.W:
W.requires_grad_(False)
for b in self.b:
b.requires_grad_(False)
return loss.item()
def dfa_config(eta=0.05, momentum=0.0):
"""SDIL configured to reproduce Direct Feedback Alignment exactly."""
return SDILConfig(eta=eta, use_residual=False, learn_A=False, learn_P=False,
pert_every=10**9, momentum=momentum)
@torch.no_grad()
def evaluate(net, test_loader):
correct, total, tot_loss = 0, 0, 0.0
for x, y in test_loader:
logits = net.logits(x)
tot_loss += F.cross_entropy(logits, y, reduction="sum").item()
correct += (logits.argmax(1) == y).sum().item()
total += y.shape[0]
return correct / total, tot_loss / total
|
