diff options
Diffstat (limited to 'experiments/conv_run.py')
| -rw-r--r-- | experiments/conv_run.py | 174 |
1 files changed, 148 insertions, 26 deletions
diff --git a/experiments/conv_run.py b/experiments/conv_run.py index 738c86d..d2c0653 100644 --- a/experiments/conv_run.py +++ b/experiments/conv_run.py @@ -12,11 +12,14 @@ import time import torch sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) -from sdil.conv import (CIFARHierarchicalFAResNet, CIFARKPResNet, - CIFARLocalResNet, CIFARSDILResNet, ConvSDILConfig, +from sdil.conv import (CIFARHierarchicalFAResNet, CIFARKPMixedTrafficResNet, + CIFARKPResNet, CIFARLocalResNet, CIFARSDILResNet, + ConvSDILConfig, conv_alignment_report, conv_apical_calibration_step, conv_hierarchical_alignment_report, conv_hierarchical_step, conv_kolen_pollack_step, + conv_kp_mixed_traffic_alignment_report, + conv_kp_mixed_traffic_step, conv_learned_hierarchical_step, conv_local_step, evaluate_conv, hierarchical_feedback_tracking_report, hierarchical_parameter_subspace_calibration) @@ -108,17 +111,25 @@ def build(args): bn_momentum=args.bn_momentum, bn_eps=args.bn_eps) if args.mode == "bp": return CIFARLocalResNet(**common), None - if args.mode in ("hfa", "lhfa", "wm", "rrm", "kp"): - network_class = ( - CIFARKPResNet if args.mode == "kp" else CIFARHierarchicalFAResNet) - net = network_class( - **common, feedback_seed=args.apical_seed, - feedback_scale=args.a_scale) + if args.mode in ("hfa", "lhfa", "wm", "rrm", "kp", "kp_traffic"): + if args.mode == "kp_traffic": + net = CIFARKPMixedTrafficResNet( + **common, feedback_seed=args.apical_seed, + feedback_scale=args.a_scale, traffic_seed=args.traffic_seed) + net.traffic_rule = args.traffic_rule + else: + network_class = ( + CIFARKPResNet if args.mode == "kp" + else CIFARHierarchicalFAResNet) + net = network_class( + **common, feedback_seed=args.apical_seed, + feedback_scale=args.a_scale) config = ConvSDILConfig( eta=args.lr, eta_output=args.output_lr, eta_A=(args.eta_A if args.mode == "lhfa" else 0.0), + eta_P=args.eta_P, momentum=args.momentum, weight_decay=args.weight_decay, - learn_A=args.mode == "lhfa", learn_P=False, + learn_A=args.mode == "lhfa", learn_P=args.mode == "kp_traffic", pert_sigma=args.pert_sigma, pert_every=args.pert_every, pert_directions=args.pert_directions, apical_calibration_mode="hierarchical_parameter_subspace") @@ -145,7 +156,9 @@ def work_report(net, mode, counters): apical_macs = getattr(net, "apical_macs_per_example", 0) normal_forward = counters["ordinary_examples"] * forward_macs warmup_forward = ((counters["predictor_warmup_examples"] - + counters["apical_warmup_examples"]) * forward_macs) + + counters["apical_warmup_examples"] + + counters["traffic_calibration_examples"] + + counters["traffic_audit_examples"]) * forward_macs) calibration_forward = counters["perturbation_forward_examples"] * forward_macs if mode == "bp": bp_reverse = 2 * counters["ordinary_examples"] * forward_macs @@ -156,7 +169,9 @@ def work_report(net, mode, counters): bp_reverse = 0 local_correlation = counters["ordinary_examples"] * forward_macs apical_inference = ((counters["ordinary_examples"] - + counters["apical_warmup_examples"]) * apical_macs) + + counters["apical_warmup_examples"] + + counters["traffic_calibration_examples"]) + * apical_macs) regression_multiplier = 2 if mode == "lhfa" else 1 apical_regression = (regression_multiplier * counters["calibration_event_examples"] @@ -168,7 +183,17 @@ def work_report(net, mode, counters): mirror_prediction = mirror_forward if mode == "rrm" else 0 mirror_correlation = mirror_forward kp_feedback_correlation = ( - counters["ordinary_examples"] * apical_macs if mode == "kp" else 0) + counters["ordinary_examples"] * apical_macs + if mode in ("kp", "kp_traffic") else 0) + if mode == "kp_traffic": + elementwise_operations = ( + counters["ordinary_examples"] + * net.mixed_elementwise_ops_per_example() + + (counters["predictor_warmup_examples"] + + counters["predictor_update_examples"]) + * net.predictor_elementwise_ops_per_example) + else: + elementwise_operations = 0 components = { "ordinary_forward_macs": normal_forward, "warmup_clean_forward_macs": warmup_forward, @@ -187,12 +212,17 @@ def work_report(net, mode, counters): "apical_macs_per_example": apical_macs, "components": components, "total_macs_estimate": sum(components.values()), + "elementwise_operations_estimate": elementwise_operations, "total_clean_forward_examples": ( counters["ordinary_examples"] + counters["predictor_warmup_examples"] - + counters["apical_warmup_examples"]), + + counters["apical_warmup_examples"] + + counters["traffic_calibration_examples"] + + counters["traffic_audit_examples"]), "total_forward_equivalent_examples": ( counters["ordinary_examples"] + counters["predictor_warmup_examples"] + counters["apical_warmup_examples"] + + counters["traffic_calibration_examples"] + + counters["traffic_audit_examples"] + counters["perturbation_forward_examples"]), "logical_batch_loss_queries": counters["logical_batch_loss_queries"], "causal_scalar_observations": counters["causal_scalar_observations"], @@ -203,18 +233,28 @@ def work_report(net, mode, counters): "multiply-accumulates in conv/linear maps; one local weight correlation " "equals one forward-weight MAC count; BP reverse is estimated as one " "weight-gradient plus one activation-gradient convolution per forward " - "convolution; elementwise nonlinearities and optimizer arithmetic excluded"), + "convolution; mixed-traffic/predictor elementwise arithmetic is reported " + "separately and is not folded into MACs"), } def run(args): if args.eval_split == "test" and args.eval_every: raise ValueError("test protocols must use --eval_every 0 (one final evaluation)") - if args.mode not in ("sdil", "lhfa") and ( + if args.mode not in ("sdil", "lhfa", "kp_traffic") and ( args.a_warmup_steps or args.learn_P): raise ValueError("apical/predictor warmup is restricted to SDIL") if args.mode == "lhfa" and args.learn_P: raise ValueError("predictor learning is not defined for learned HFA") + if args.mode == "kp_traffic": + if not args.learn_P or args.predictor_warmup_steps < 1: + raise ValueError("mixed-traffic KP requires predictor warmup") + if args.predictor_every < 1: + raise ValueError("mixed-traffic KP requires a predictor cadence") + if args.traffic_calibration_examples < 1 or args.traffic_ratio <= 0: + raise ValueError("invalid mixed-traffic calibration") + elif args.predictor_every: + raise ValueError("predictor cadence is restricted to mixed-traffic KP") if args.mode not in ("wm", "rrm") and args.mirror_warmup_steps: raise ValueError("mirror warmup is restricted to weight mirror mode") if args.mirror_every < 1 or args.mirror_batch_size < 1: @@ -262,12 +302,17 @@ def run(args): "mirror_conv_examples": 0, "mirror_readout_examples": 0, "mirror_events": 0, + "traffic_calibration_examples": 0, + "traffic_audit_examples": 0, + "predictor_update_examples": 0, } log = { "schema_version": 1, "protocol_family": "oral_a_cifar_local_resnet_development", "calibration_metric_space": ( None if config is None or args.mode == "hfa" else + "reciprocal_local_activity_products_with_mixed_apical_traffic" + if args.mode == "kp_traffic" else "hierarchical_feedback_parameters" if args.mode == "lhfa" else "reciprocal_local_activity_products" if args.mode == "kp" else { "wm": "local_parent_child_response", @@ -302,7 +347,8 @@ def run(args): if args.mode == "hfa" else 0), "adaptive_feedback_parameters": ( getattr(net, "n_fixed_feedback_parameters", 0) - if args.mode in ("lhfa", "wm", "rrm", "kp") else 0), + if args.mode in ("lhfa", "wm", "rrm", "kp", "kp_traffic") + else 0), }, "epochs": [], } @@ -313,6 +359,33 @@ def run(args): apical_warmup_wall = 0.0 mirror_warmup_wall = 0.0 loader_state = train.g.get_state().clone() + traffic_calibration_batch = None + if args.mode == "kp_traffic": + count = min(args.traffic_calibration_examples, train.x.shape[0]) + if count != args.traffic_calibration_examples: + raise ValueError("traffic calibration prefix is unavailable") + calibration_x = train.x[:count] + calibration_y = train.y[:count] + traffic_calibration_batch = (calibration_x, calibration_y) + calibration_forward = net.forward( + calibration_x, return_cache=True, training=True, + update_stats=False) + calibration_output_error = ( + torch.softmax(calibration_forward["logits"], dim=1) + - torch.nn.functional.one_hot( + calibration_y, net.n_classes).to( + calibration_forward["logits"].dtype)) + calibration_instruction = net.hierarchical_teaching( + calibration_output_error, calibration_forward) + log["traffic_calibration"] = net.calibrate_traffic_gain( + calibration_instruction, calibration_forward["hiddens"], + args.traffic_ratio) + log["traffic_calibration"].update({ + "examples": count, + "data_source": "first unaugmented development-training examples", + "uses_validation_endpoint": False, + }) + counters["traffic_calibration_examples"] += count if args.mode in ("wm", "rrm") and args.mirror_warmup_steps: sync(args.device) mirror_start = time.time() @@ -342,6 +415,7 @@ def run(args): sync(args.device) warmup_start = time.time() iterator = iter(train) + predictor_metrics = [] for _ in range(args.predictor_warmup_steps): try: x, _ = next(iterator) @@ -349,10 +423,30 @@ def run(args): iterator = iter(train) x, _ = next(iterator) forward = net.forward(x, training=True, update_stats=False) - net.predictor_step( - forward["hiddens"], config.eta_P, config.nuisance_scale) + predictor_metric = (net.predictor_step( + forward["hiddens"], config.eta_P) + if args.mode == "kp_traffic" else net.predictor_step( + forward["hiddens"], config.eta_P, + config.nuisance_scale)) + predictor_metrics.append(predictor_metric) counters["predictor_warmup_examples"] += x.shape[0] train.g.set_state(loader_state) + log["predictor_warmup"] = { + "steps": args.predictor_warmup_steps, + "first_mse": predictor_metrics[0], + "mean_mse": sum(predictor_metrics) / len(predictor_metrics), + "last_mse": predictor_metrics[-1], + "instruction_present": False, + } + if args.mode == "kp_traffic": + audit_x, _ = traffic_calibration_batch + audit_forward = net.forward( + audit_x, training=True, update_stats=False) + log["predictor_warmup"][ + "post_warmup_traffic_residual_rms_ratio"] = ( + net.predictor_traffic_residual_rms_ratio( + audit_forward["hiddens"])) + counters["traffic_audit_examples"] += audit_x.shape[0] sync(args.device) predictor_warmup_wall = time.time() - warmup_start @@ -426,6 +520,7 @@ def run(args): examples = 0 calibration_metrics = [] mirror_metrics = [] + signal_metrics = [] for x, y in train: batch = x.shape[0] if args.mode == "bp": @@ -476,6 +571,17 @@ def run(args): result = conv_kolen_pollack_step(net, x, y, config) loss = result["loss"] did_perturb = False + elif args.mode == "kp_traffic": + result = conv_kp_mixed_traffic_step( + net, x, y, config, step, args.traffic_rule, + args.predictor_every) + loss = result["loss"] + did_perturb = False + signal_metrics.append({key: result[key] for key in ( + "teaching_rms", "instruction_rms", "raw_apical_rms", + "innovation_rms", "traffic_rms")}) + if result["did_predictor_update"]: + counters["predictor_update_examples"] += batch else: result = conv_local_step( net, x, y, config, step, generator=perturb_generator) @@ -521,7 +627,12 @@ def run(args): key: sum(value[key] for value in mirror_metrics) / len(mirror_metrics) for key in mirror_metrics[0] } - if args.mode == "kp": + if signal_metrics: + record["mixed_apical"] = { + key: sum(value[key] for value in signal_metrics) + / len(signal_metrics) for key in signal_metrics[0] + } + if args.mode in ("kp", "kp_traffic"): tracking = hierarchical_feedback_tracking_report(net) record["feedback_tracking"] = { "mean_feedback_forward_cosine": tracking[ @@ -569,11 +680,15 @@ def run(args): probe = min(args.alignment_probe, train.x.shape[0]) sync(args.device) diagnostic_start = time.time() - diagnostics = (conv_hierarchical_alignment_report( - net, train.x[:probe], train.y[:probe]) - if args.mode in ("hfa", "lhfa", "wm", "rrm", "kp") else - conv_alignment_report( - net, train.x[:probe], train.y[:probe], config)) + if args.mode == "kp_traffic": + diagnostics = conv_kp_mixed_traffic_alignment_report( + net, train.x[:probe], train.y[:probe], args.traffic_rule) + elif args.mode in ("hfa", "lhfa", "wm", "rrm", "kp"): + diagnostics = conv_hierarchical_alignment_report( + net, train.x[:probe], train.y[:probe]) + else: + diagnostics = conv_alignment_report( + net, train.x[:probe], train.y[:probe], config) sync(args.device) diagnostics["wall_s"] = time.time() - diagnostic_start values = diagnostics["teaching_negative_gradient_cosine"] @@ -620,8 +735,8 @@ def parse_args(): parser = argparse.ArgumentParser() parser.add_argument( "--mode", choices=( - "bp", "dfa", "hfa", "lhfa", "wm", "rrm", "kp", "sdil", - "nodepert"), + "bp", "dfa", "hfa", "lhfa", "wm", "rrm", "kp", + "kp_traffic", "sdil", "nodepert"), required=True) parser.add_argument("--out", required=True) parser.add_argument("--device", default="cpu") @@ -673,6 +788,13 @@ def parse_args(): choices=("unit_targets", "channel_subspace", "vectorizer_subspace"), default="unit_targets") parser.add_argument("--predictor_warmup_steps", type=int, default=0) + parser.add_argument("--predictor_every", type=int, default=0) + parser.add_argument("--traffic_rule", + choices=("raw", "matched", "innovation"), + default="innovation") + parser.add_argument("--traffic_seed", type=int, default=4000) + parser.add_argument("--traffic_ratio", type=float, default=4.0) + parser.add_argument("--traffic_calibration_examples", type=int, default=64) parser.add_argument("--a_warmup_steps", type=int, default=0) parser.add_argument("--mirror_warmup_steps", type=int, default=0) parser.add_argument("--mirror_every", type=int, default=16) |