experiments: audit calibration cost and isolate timing

1 files changed, 16 insertions, 0 deletions

diff --git a/experiments/protocol_smoke.py b/experiments/protocol_smoke.py
index 36da996..551987c 100644
--- a/experiments/protocol_smoke.py
+++ b/experiments/protocol_smoke.py
@@ -6,6 +6,8 @@ import torch
 
 sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 from sdil.data import get_dataset_splits
+from sdil.core import SDILConfig, SDILNet
+from experiments.run import calibration_work_per_event
 
 
 def loader_labels(loader):
@@ -30,8 +32,22 @@ def main():
     assert sum(y.numel() for _, y in train) == 59000
     assert sum(y.numel() for _, y in validation) == 1000
     assert sum(y.numel() for _, y in test) == 10000
+
+    net = SDILNet([10, 8, 8, 3], device="cpu")
+    simultaneous = calibration_work_per_event(
+        net, SDILConfig(pert_ndirs=2, pert_mode="simultaneous"))
+    layerwise = calibration_work_per_event(
+        net, SDILConfig(pert_ndirs=2, pert_mode="layerwise"))
+    assert simultaneous == {
+        "batch_loss_evaluations": 4,
+        "forward_equivalent_batches": 5.0,
+        "perturbation_batch_expansion": 4,
+    }
+    assert layerwise["batch_loss_evaluations"] == 9
+    assert abs(layerwise["forward_equivalent_batches"] - 11.0 / 3.0) < 1e-12
     print("validation split hash:", metadata["validation_index_sha256"])
     print("train/validation/test: 59000/1000/10000; stratification exact")
+    print("simultaneous/layerwise calibration cost accounting: exact")
     print("ALL PROTOCOL SMOKE CHECKS PASSED")