protocol: freeze post-failure Oral-A v2 funnel

1 files changed, 12 insertions, 0 deletions

diff --git a/experiments/conv_run.py b/experiments/conv_run.py
index cb99545..a33005d 100644
--- a/experiments/conv_run.py
+++ b/experiments/conv_run.py
@@ -212,6 +212,11 @@ def run(args):
     log = {
         "schema_version": 1,
         "protocol_family": "oral_a_cifar_local_resnet_development",
+        "calibration_metric_space": (
+            None if config is None else
+            ("channel_basis_moments"
+             if config.apical_calibration_mode == "channel_subspace"
+             else "full_hidden_field")),
         "args": vars(args),
         "provenance": provenance(),
         "split": split,
@@ -284,8 +289,15 @@ def run(args):
             counters["per_example_loss_terms"] += 2 * config.pert_directions * batch
             counters["perturbation_events"] += 1
         train.g.set_state(loader_state)
+        warmup_mean = {
+            key: sum(metric[key] for metric in warmup_metrics)
+            / len(warmup_metrics)
+            for key in warmup_metrics[0]
+        }
         log["apical_warmup"] = {
             "steps": args.a_warmup_steps,
+            "first": warmup_metrics[0],
+            "mean": warmup_mean,
             "last": warmup_metrics[-1],
         }
         sync(args.device)