oral-a: add exact local convolution eligibilities

1 files changed, 123 insertions, 0 deletions

diff --git a/experiments/conv_local_smoke.py b/experiments/conv_local_smoke.py
new file mode 100644
index 0000000..35073a6
--- /dev/null
+++ b/experiments/conv_local_smoke.py
@@ -0,0 +1,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()