chore: capture initial SDIL project state

1 files changed, 58 insertions, 0 deletions

diff --git a/sdil/baselines.py b/sdil/baselines.py
new file mode 100644
index 0000000..f809f83
--- /dev/null
+++ b/sdil/baselines.py
@@ -0,0 +1,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