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
|
#!/usr/bin/env python3
"""Prove the convolutional local eligibility matches exact BP when instructed."""
import os
import sys
import torch
import torch.nn.functional as F
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from sdil.conv import CIFARLocalResNet
def architecture_checks():
expected = {
8: (7, 74810),
20: (19, 268346),
32: (31, 461882),
56: (55, 848954),
}
for depth, (hidden, parameters) in expected.items():
net = CIFARLocalResNet(depth=depth)
assert net.n_hidden == hidden
assert net.n_forward_parameters == parameters
assert len(net.blocks) == 3 * ((depth - 2) // 6)
assert len(net.W) + 1 == depth
for depth in (7, 9, 21):
try:
CIFARLocalResNet(depth=depth)
except ValueError:
pass
else:
raise AssertionError(f"invalid depth {depth} was accepted")
x = torch.arange(2 * 3 * 8 * 8, dtype=torch.float32).reshape(2, 3, 8, 8)
shortcut = CIFARLocalResNet._option_a_shortcut(x, 6, 2)
assert tuple(shortcut.shape) == (2, 6, 4, 4)
assert torch.count_nonzero(shortcut[:, 0]) == 0
assert torch.count_nonzero(shortcut[:, -2:]) == 0
assert torch.equal(shortcut[:, 1:4], x[:, :, ::2, ::2])
def exact_local_gradient_check():
torch.manual_seed(123)
batch = 3
x = torch.randn(batch, 3, 32, 32)
y = torch.tensor([0, 3, 8])
local = CIFARLocalResNet(depth=8, base_width=4, seed=19)
bp = CIFARLocalResNet(depth=8, base_width=4, seed=19)
parameters = local.W + [local.W_out, local.b_out]
for parameter in parameters:
parameter.requires_grad_(True)
forward = local.forward(x, return_cache=True)
for hidden in forward["hiddens"]:
hidden.retain_grad()
loss = F.cross_entropy(forward["logits"], y)
loss.backward()
# .backward() differentiated a batch-mean loss. Multiplying hidden grads
# by B recovers the per-example convention consumed by the local rule.
teaching = [-batch * hidden.grad for hidden in forward["hiddens"]]
output_error = (torch.softmax(forward["logits"].detach(), dim=1)
- F.one_hot(y, local.n_classes))
directions, out_direction, bias_direction = local.local_ascent_directions(
teaching, output_error, forward)
relative_errors = []
for direction, parameter in zip(directions, local.W):
absolute = (direction + parameter.grad).abs().max()
scale = parameter.grad.abs().max().clamp_min(1e-12)
relative_errors.append(float(absolute / scale))
output_abs = float((out_direction + local.W_out.grad).abs().max())
bias_abs = float((bias_direction + local.b_out.grad).abs().max())
assert max(relative_errors) < 3e-5
assert output_abs < 2e-6 and bias_abs < 2e-6
for parameter in parameters:
parameter.requires_grad_(False)
eta = 0.017
local.apply_ascent(directions, out_direction, bias_direction, eta)
bp_loss = bp.bp_step(x, y, eta)
assert abs(float(loss.detach()) - bp_loss) < 1e-7
parameter_differences = [
float((left - right).abs().max())
for left, right in zip(
local.W + [local.W_out, local.b_out],
bp.W + [bp.W_out, bp.b_out])]
assert max(parameter_differences) < 2e-7
return {
"max_relative_local_gradient_error": max(relative_errors),
"output_absolute_error": output_abs,
"post_update_parameter_max_error": max(parameter_differences),
}
def perturbation_checks():
net = CIFARLocalResNet(depth=8, base_width=4, seed=3)
x = torch.randn(2, 3, 32, 32)
clean = net.forward(x)
perturbations = [torch.zeros_like(hidden) for hidden in clean["hiddens"]]
perturbed = net.forward(x, perturbations=perturbations)
assert torch.equal(clean["logits"], perturbed["logits"])
perturbations[0] = torch.ones_like(perturbations[0]) * 0.01
changed = net.forward(x, perturbations=perturbations)
assert not torch.equal(clean["logits"], changed["logits"])
try:
net.forward(x, perturbations=perturbations[:-1])
except ValueError:
pass
else:
raise AssertionError("short perturbation list was accepted")
def main():
architecture_checks()
perturbation_checks()
report = exact_local_gradient_check()
print(report)
print("ALL CONVOLUTIONAL LOCAL-ELIGIBILITY CHECKS PASSED")
if __name__ == "__main__":
main()
|
