#!/usr/bin/env python3 """Run the frozen no-KP layerwise causal-bootstrap capture screen.""" import argparse import json import math import os import subprocess import sys import time import torch sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) from sdil.conv import (CIFARHierarchicalFAResNet, conv_hierarchical_alignment_report, layerwise_causal_bootstrap_sweep) from sdil.data import DATA_DIR, get_cifar_image_splits ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) def provenance(): def run(command): return subprocess.run( command, cwd=ROOT, check=True, capture_output=True, text=True).stdout.strip() return { "git_commit": run(["git", "rev-parse", "HEAD"]), "git_tracked_dirty": bool(run( ["git", "status", "--porcelain", "--untracked-files=no"])), } def summarize_alignment(report): values = report["teaching_negative_gradient_cosine"] early = max(1, len(values) // 3) ratios = report["feedback_forward_norm_ratio"] cosines = report["feedback_forward_cosine"] return { "per_layer": values, "early_third_alignment": sum(values[:early]) / early, "all_layer_alignment": sum(values) / len(values), "mean_feedback_forward_cosine": sum(cosines) / len(cosines), "min_feedback_forward_norm_ratio": min(ratios), "max_feedback_forward_norm_ratio": max(ratios), "feedback_forward_cosine": cosines, "feedback_forward_norm_ratio": ratios, } def forward_state(net): return [value.clone() for value in ( net.W + net.gamma + net.beta + net.running_mean + net.running_var + net.mW + net.mgamma + net.mbeta + [net.W_out, net.b_out, net.mW_out, net.mb_out])] def main(): parser = argparse.ArgumentParser() parser.add_argument("--device", default="cuda") parser.add_argument("--data_dir", default=DATA_DIR) parser.add_argument("--out", default="results/oral_a_v5_calibration/result.json") args = parser.parse_args() settings = { "depth": 20, "width": 16, "seed": 0, "loader_seed": 0, "batch_size": 128, "train_limit": 10000, "val_examples": 5000, "split_seed": 2027, "normalization": "batchnorm", "residual_scale": 1.0, "feedback_scale": 1.0, "sigma": 0.01, "eta_A": 0.1, "perturb_seed": 5000, "sweeps": 20, "alignment_probe": 64, "calibration_augmentation": False, } torch.manual_seed(settings["seed"]) if str(args.device).startswith("cuda"): if not torch.cuda.is_available(): raise RuntimeError("CUDA requested but unavailable") torch.cuda.manual_seed_all(settings["seed"]) torch.cuda.reset_peak_memory_stats(torch.device(args.device)) train, _, _, input_shape, n_out, split = get_cifar_image_splits( batch_size=settings["batch_size"], data_dir=args.data_dir, device=args.device, train_limit=settings["train_limit"], val_examples=settings["val_examples"], split_seed=settings["split_seed"], loader_seed=settings["loader_seed"], augment_train=False) if input_shape != (3, 32, 32) or n_out != 10: raise AssertionError("unexpected CIFAR dimensions") net = CIFARHierarchicalFAResNet( depth=settings["depth"], base_width=settings["width"], n_classes=10, device=args.device, seed=settings["seed"], residual_scale=settings["residual_scale"], normalization=settings["normalization"], feedback_scale=settings["feedback_scale"]) audit_x = train.x[:settings["alignment_probe"]] audit_y = train.y[:settings["alignment_probe"]] fixed = summarize_alignment( conv_hierarchical_alignment_report(net, audit_x, audit_y)) state_before = forward_state(net) generator = torch.Generator(device=torch.device(args.device)).manual_seed( settings["perturb_seed"]) if str(args.device).startswith("cuda"): torch.cuda.synchronize(torch.device(args.device)) start = time.time() sweeps = [] for sweep_index in range(settings["sweeps"]): start_index = sweep_index * settings["batch_size"] stop_index = start_index + settings["batch_size"] metric = layerwise_causal_bootstrap_sweep( net, train.x[start_index:stop_index], train.y[start_index:stop_index], sigma=settings["sigma"], eta=settings["eta_A"], generator=generator) sweeps.append({ key: value for key, value in metric.items() if key != "edges"}) print(json.dumps({"sweep": sweep_index + 1, **sweeps[-1]}), flush=True) if str(args.device).startswith("cuda"): torch.cuda.synchronize(torch.device(args.device)) wall_seconds = time.time() - start state_after = forward_state(net) forward_state_max_difference = max( float((before - after).abs().max()) for before, after in zip(state_before, state_after)) learned = summarize_alignment( conv_hierarchical_alignment_report(net, audit_x, audit_y)) total_events = sum(value["events"] for value in sweeps) total_queries = sum(value["logical_batch_loss_queries"] for value in sweeps) total_observations = sum( value["per_example_causal_observations"] for value in sweeps) batch = settings["batch_size"] clean_forward_examples = settings["sweeps"] * batch perturbation_forward_examples = 2 * total_events * batch teaching_macs = total_events * batch * net.apical_macs_per_example work = { "edge_events": total_events, "logical_batch_loss_queries": total_queries, "per_example_causal_observations": total_observations, "per_example_cross_entropy_terms": 2 * total_observations, "clean_forward_examples": clean_forward_examples, "perturbation_forward_examples": perturbation_forward_examples, "forward_macs": ((clean_forward_examples + perturbation_forward_examples) * net.forward_macs_per_example), "hierarchical_teaching_and_local_correlation_macs_estimate": teaching_macs, } work["total_macs_estimate"] = ( work["forward_macs"] + work["hierarchical_teaching_and_local_correlation_macs_estimate"]) finite_values = [ fixed["early_third_alignment"], fixed["all_layer_alignment"], learned["early_third_alignment"], learned["all_layer_alignment"], learned["min_feedback_forward_norm_ratio"], learned["max_feedback_forward_norm_ratio"], ] + [value[key] for value in sweeps for key in ( "mean_field_prediction_target_cosine", "mean_parameter_update_rms", "max_parameter_update_rms")] output = { "schema_version": 1, "protocol": "oral_a_v5_layerwise_causal_bootstrap_capture_v1", "settings": settings, "provenance": provenance(), "split": split, "architecture": { "family": "CIFAR 6n+2 ResNet, option-A shortcuts", "forward_parameters": net.n_forward_parameters, "adaptive_feedback_parameters": net.n_fixed_feedback_parameters, "forward_macs_per_example": net.forward_macs_per_example, "feedback_macs_per_example": net.apical_macs_per_example, }, "method_audit": { "forward_weight_reads_in_feedback_update": 0, "reverse_mode_learning_operations": 0, "causal_query_normalization_state": "evaluation_running_statistics", "ordinary_task_normalization_state": "not_run_forward_frozen", "forward_state_max_absolute_difference": ( forward_state_max_difference), }, "fixed_hfa": fixed, "learned_lcb": learned, "sweeps": sweeps, "work": work, "wall_seconds": wall_seconds, "finite": all(math.isfinite(value) for value in finite_values), "test_examples_touched": 0, "validation_endpoints_observed": 0, "hardware": { "device": str(args.device), "torch_version": torch.__version__, "cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"), "cuda_device_name": (torch.cuda.get_device_name(torch.device(args.device)) if str(args.device).startswith("cuda") else None), "peak_memory_allocated_bytes": ( torch.cuda.max_memory_allocated(torch.device(args.device)) if str(args.device).startswith("cuda") else None), }, } os.makedirs(os.path.dirname(os.path.abspath(args.out)), exist_ok=True) with open(args.out, "w") as handle: json.dump(output, handle, indent=2, sort_keys=True) handle.write("\n") print(json.dumps({ "fixed_hfa": fixed, "learned_lcb": learned, "work": work, "finite": output["finite"], "wall_seconds": wall_seconds, "out": args.out, }, indent=2), flush=True) if __name__ == "__main__": main()