1 files changed, 94 insertions, 0 deletions
diff --git a/experiments/resnet_crossover_smoke.py b/experiments/resnet_crossover_smoke.py
new file mode 100644
index 0000000..fdc92fa
--- /dev/null
+++ b/experiments/resnet_crossover_smoke.py
@@ -0,0 +1,94 @@
+#!/usr/bin/env python3
+"""Deterministic equation audits for matched ResNet crossover adapters."""
+import json
+import os
+import sys
+
+import torch
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+from sdil.conv import CIFARLocalResNet
+from sdil.conv_crossover import CIFARDualPropResNet
+
+
+def relative_error(actual, expected):
+    numerator = torch.linalg.vector_norm(actual - expected)
+    denominator = torch.linalg.vector_norm(expected).clamp_min(1e-30)
+    return float(numerator / denominator)
+
+
+def audit_dualprop():
+    common = dict(
+        depth=8, base_width=2, seed=71, normalization="batchnorm",
+        residual_scale=1.0, dtype=torch.float64)
+    reference = CIFARLocalResNet(**common)
+    net = CIFARDualPropResNet(
+        **common, alpha=0.0, dp_beta=0.1, inference_passes=2)
+    generator = torch.Generator().manual_seed(72)
+    image = torch.randn(3, 3, 32, 32, generator=generator,
+                        dtype=torch.float64)
+    labels = torch.tensor([0, 3, 7])
+    one_hot = torch.nn.functional.one_hot(labels, 10).to(torch.float64)
+    with torch.no_grad():
+        reference_output = reference.forward(
+            image, training=True, update_stats=False)["logits"]
+        net_output = net.forward(
+            image, training=True, update_stats=False)["logits"]
+        clean = net.forward(
+            image, return_cache=True, training=True, update_stats=False)
+        plus, minus = net.infer_dual_states(image, one_hot, clean)
+    forward_error = float(torch.max(torch.abs(
+        reference_output - net_output)))
+
+    parameters = (
+        net.W + net.gamma + net.beta + [net.W_out, net.b_out])
+    for parameter in parameters:
+        parameter.requires_grad_(True)
+    alpha = net.dp_alpha
+    beta = net.dp_beta
+    states = [
+        (alpha * positive + (1.0 - alpha) * negative).detach()
+        for positive, negative in zip(plus[:-1], minus[:-1])]
+    deltas = [
+        ((positive - negative) / beta).detach()
+        for positive, negative in zip(plus, minus)]
+    objective = image.new_zeros(())
+    for index in range(net.n_hidden):
+        prediction, _ = net._node_prediction(index, states, image)
+        objective -= torch.sum(deltas[index] * prediction) / image.shape[0]
+    features = states[-1].mean(dim=(2, 3))
+    output_prediction = features @ net.W_out.t() + net.b_out
+    objective -= torch.sum(deltas[-1] * output_prediction) / image.shape[0]
+    gradients = torch.autograd.grad(objective, parameters)
+    (directions, gamma_directions, beta_directions,
+     output_weight, output_bias) = net.dualprop_ascent_directions(
+         image, plus, minus)
+    actual = (
+        directions + gamma_directions + beta_directions
+        + [output_weight, output_bias])
+    errors = [
+        relative_error(direction, -gradient)
+        for direction, gradient in zip(actual, gradients)]
+    for parameter in parameters:
+        parameter.requires_grad_(False)
+    if forward_error >= 1e-12 or max(errors) >= 2e-12:
+        raise AssertionError({
+            "forward_error": forward_error,
+            "direction_errors": errors,
+        })
+    return {
+        "matched_forward_max_absolute_error": forward_error,
+        "contrastive_direction_max_relative_error": max(errors),
+        "num_audited_parameter_tensors": len(errors),
+        "uses_symmetric_forward_edge_transposes": True,
+        "uses_reverse_task_loss_graph": False,
+    }
+
+
+def main():
+    print(json.dumps({"dualprop": audit_dualprop()},
+                     indent=2, sort_keys=True))
+
+
+if __name__ == "__main__":
+    main()