diff options
Diffstat (limited to 'experiments')
| -rw-r--r-- | experiments/conv_local_smoke.py | 129 | ||||
| -rw-r--r-- | experiments/conv_run.py | 32 |
2 files changed, 150 insertions, 11 deletions
diff --git a/experiments/conv_local_smoke.py b/experiments/conv_local_smoke.py index 4a0f47d..d2428eb 100644 --- a/experiments/conv_local_smoke.py +++ b/experiments/conv_local_smoke.py @@ -24,6 +24,11 @@ def architecture_checks(): assert net.n_forward_parameters == parameters assert len(net.blocks) == 3 * ((depth - 2) // 6) assert len(net.W) + 1 == depth + batchnorm_parameters = {8: 75290, 20: 269722, 32: 464154, 56: 853018} + for depth, parameters in batchnorm_parameters.items(): + net = CIFARLocalResNet( + depth=depth, normalization="batchnorm", residual_scale=1.0) + assert net.n_forward_parameters == parameters for depth in (7, 9, 21): try: CIFARLocalResNet(depth=depth) @@ -62,8 +67,10 @@ def exact_local_gradient_check(): teaching = [-batch * hidden.grad for hidden in forward["hiddens"]] output_error = (torch.softmax(forward["logits"].detach(), dim=1) - F.one_hot(y, local.n_classes)) - directions, out_direction, bias_direction = local.local_ascent_directions( + (directions, gamma_directions, beta_directions, + out_direction, bias_direction) = local.local_ascent_directions( teaching, output_error, forward) + assert gamma_directions == beta_directions == [] relative_errors = [] for direction, parameter in zip(directions, local.W): @@ -94,6 +101,78 @@ def exact_local_gradient_check(): } +def exact_batchnorm_local_gradient_check(): + torch.manual_seed(31) + batch = 4 + x = torch.randn(batch, 3, 32, 32) + y = torch.tensor([0, 1, 2, 3]) + common = dict( + depth=8, base_width=2, seed=29, + normalization="batchnorm", residual_scale=1.0) + local = CIFARLocalResNet(**common) + bp = CIFARLocalResNet(**common) + parameters = local.W + local.gamma + local.beta + [local.W_out, local.b_out] + for parameter in parameters: + parameter.requires_grad_(True) + forward = local.forward( + x, return_cache=True, training=True, update_stats=True) + for hidden in forward["hiddens"]: + hidden.retain_grad() + loss = F.cross_entropy(forward["logits"], y) + loss.backward() + teaching = [-batch * hidden.grad for hidden in forward["hiddens"]] + output_error = (torch.softmax(forward["logits"].detach(), dim=1) + - F.one_hot(y, 10)) + (directions, gamma_directions, beta_directions, + out_direction, bias_direction) = local.local_ascent_directions( + teaching, output_error, forward) + groups = ( + (directions, local.W), + (gamma_directions, local.gamma), + (beta_directions, local.beta), + ) + relative_errors = [] + for direction_group, parameter_group in groups: + for direction, parameter in zip(direction_group, parameter_group): + absolute = (direction + parameter.grad).abs().max() + relative_errors.append(float( + absolute / parameter.grad.abs().max().clamp_min(1e-12))) + assert max(relative_errors) < 3e-5 + for parameter in parameters: + parameter.requires_grad_(False) + + eta = 0.013 + local.apply_ascent( + directions, out_direction, bias_direction, eta, + gamma_directions=gamma_directions, beta_directions=beta_directions) + bp.bp_step(x, y, eta) + parameter_differences = [ + float((left - right).abs().max()) + for left, right in zip( + local.W + local.gamma + local.beta + [local.W_out, local.b_out], + bp.W + bp.gamma + bp.beta + [bp.W_out, bp.b_out])] + running_differences = [ + float((left - right).abs().max()) + for left, right in zip( + local.running_mean + local.running_var, + bp.running_mean + bp.running_var)] + assert max(parameter_differences) < 2e-7 + assert max(running_differences) == 0.0 + + running_before = [value.clone() for value in local.running_mean + local.running_var] + clean = local.forward(x, training=True, update_stats=False) + simultaneous_conv_node_perturbation( + local, x, y, clean, sigma=1e-3, n_directions=1, + generator=torch.Generator(device="cpu").manual_seed(9)) + assert all(torch.equal(before, after) for before, after in zip( + running_before, local.running_mean + local.running_var)) + assert torch.equal(local.logits(x), local.logits(x)) + return { + "batchnorm_max_relative_local_gradient_error": max(relative_errors), + "batchnorm_post_update_parameter_max_error": max(parameter_differences), + } + + def perturbation_checks(): net = CIFARLocalResNet(depth=8, base_width=4, seed=3) x = torch.randn(2, 3, 32, 32) @@ -132,7 +211,9 @@ def perturbation_estimator_check(): net, x, y, clean, sigma=1e-6, n_directions=1, generator=generator, return_diagnostics=True) directions = diagnostics["directions"][0] - finite_difference = diagnostics["directional_derivatives"][0] + derivative = diagnostics["directional_derivatives"][0] + assert derivative["coupling"] == "per_example_objective" + finite_difference = derivative["scaled_directional"] exact = sum( (batch * hidden.grad * direction).flatten(1).sum(dim=1) for hidden, direction in zip(clean["hiddens"], directions)) @@ -143,7 +224,46 @@ def perturbation_estimator_check(): assert torch.equal(target, expected) for parameter in parameters: parameter.requires_grad_(False) - return float(relative.max()) + no_norm_relative = float(relative.max()) + + batch = 3 + batchnorm = CIFARSDILResNet( + depth=8, base_width=2, seed=14, dtype=torch.float64, + normalization="batchnorm", residual_scale=1.0) + xb = torch.randn(batch, 3, 32, 32, dtype=torch.float64) + yb = torch.tensor([1, 4, 9]) + parameters = (batchnorm.W + batchnorm.gamma + batchnorm.beta + + [batchnorm.W_out, batchnorm.b_out]) + for parameter in parameters: + parameter.requires_grad_(True) + clean = batchnorm.forward(xb, training=True, update_stats=False) + for hidden in clean["hiddens"]: + hidden.retain_grad() + F.cross_entropy(clean["logits"], yb).backward() + targets, diagnostics = simultaneous_conv_node_perturbation( + batchnorm, xb, yb, clean, sigma=1e-6, n_directions=1, + generator=torch.Generator(device="cpu").manual_seed(101), + return_diagnostics=True) + directions = diagnostics["directions"][0] + derivative = diagnostics["directional_derivatives"][0] + assert derivative["coupling"] == "batch_objective" + scaled = derivative["scaled_directional"] + exact_sum_directional = sum( + float((batch * hidden.grad * direction).sum()) + for hidden, direction in zip(clean["hiddens"], directions)) + batchnorm_relative = abs(float(scaled[0]) - exact_sum_directional) / max( + abs(exact_sum_directional), 1e-12) + assert batchnorm_relative < 2e-6 + assert torch.equal(scaled, scaled[:1].expand_as(scaled)) + for target, direction in zip(targets, directions): + expected = -scaled[:, None, None, None] * direction + assert torch.equal(target, expected) + for parameter in parameters: + parameter.requires_grad_(False) + return { + "perturbation_jvp_max_relative_error": no_norm_relative, + "batchnorm_batch_objective_jvp_relative_error": batchnorm_relative, + } def apical_learning_checks(): @@ -197,7 +317,8 @@ def main(): architecture_checks() perturbation_checks() report = exact_local_gradient_check() - report["perturbation_jvp_max_relative_error"] = perturbation_estimator_check() + report.update(exact_batchnorm_local_gradient_check()) + report.update(perturbation_estimator_check()) report.update(apical_learning_checks()) print(report) print("ALL CONVOLUTIONAL LOCAL-ELIGIBILITY CHECKS PASSED") diff --git a/experiments/conv_run.py b/experiments/conv_run.py index 48eebd1..dd3f6c5 100644 --- a/experiments/conv_run.py +++ b/experiments/conv_run.py @@ -91,10 +91,14 @@ def scheduled_lr(base, epoch, args): def build(args): + residual_scale = args.residual_scale + if residual_scale is None and args.normalization == "batchnorm": + residual_scale = 1.0 common = dict( depth=args.depth, base_width=args.width, n_classes=10, device=args.device, seed=args.seed, weight_scale=args.weight_scale, - residual_scale=args.residual_scale) + residual_scale=residual_scale, normalization=args.normalization, + bn_momentum=args.bn_momentum, bn_eps=args.bn_eps) if args.mode == "bp": return CIFARLocalResNet(**common), None net = CIFARSDILResNet( @@ -151,7 +155,8 @@ def work_report(net, mode, counters): + counters["apical_warmup_examples"] + counters["perturbation_forward_examples"]), "logical_batch_loss_queries": counters["logical_batch_loss_queries"], - "scalar_loss_evaluations": counters["scalar_loss_evaluations"], + "causal_scalar_observations": counters["causal_scalar_observations"], + "per_example_cross_entropy_terms": counters["per_example_loss_terms"], "definition": ( "multiply-accumulates in conv/linear maps; one local weight correlation " "equals one forward-weight MAC count; BP reverse is estimated as one " @@ -198,7 +203,8 @@ def run(args): "perturbation_forward_examples": 0, "calibration_event_examples": 0, "logical_batch_loss_queries": 0, - "scalar_loss_evaluations": 0, + "causal_scalar_observations": 0, + "per_example_loss_terms": 0, "perturbation_events": 0, } log = { @@ -208,7 +214,10 @@ def run(args): "provenance": provenance(), "split": split, "architecture": { - "family": "normalization-free CIFAR 6n+2 ResNet, option-A shortcuts", + "family": "CIFAR 6n+2 ResNet, option-A shortcuts", + "normalization": net.normalization, + "bn_momentum": net.bn_momentum if net.normalization == "batchnorm" else None, + "bn_eps": net.bn_eps if net.normalization == "batchnorm" else None, "depth": net.depth, "blocks_per_stage": net.blocks_per_stage, "base_width": net.base_width, @@ -231,7 +240,7 @@ def run(args): except StopIteration: iterator = iter(train) x, _ = next(iterator) - forward = net.forward(x) + forward = net.forward(x, training=True, update_stats=False) net.predictor_step( forward["hiddens"], config.eta_P, config.nuisance_scale) counters["predictor_warmup_examples"] += x.shape[0] @@ -255,7 +264,9 @@ def run(args): 2 * config.pert_directions * batch) counters["calibration_event_examples"] += batch counters["logical_batch_loss_queries"] += 2 * config.pert_directions - counters["scalar_loss_evaluations"] += 2 * config.pert_directions * batch + counters["causal_scalar_observations"] += 2 * config.pert_directions * ( + 1 if net.normalization == "batchnorm" else batch) + counters["per_example_loss_terms"] += 2 * config.pert_directions * batch counters["perturbation_events"] += 1 train.g.set_state(loader_state) log["apical_warmup"] = { @@ -304,7 +315,10 @@ def run(args): 2 * config.pert_directions * batch) counters["calibration_event_examples"] += batch counters["logical_batch_loss_queries"] += 2 * config.pert_directions - counters["scalar_loss_evaluations"] += ( + counters["causal_scalar_observations"] += ( + 2 * config.pert_directions + * (1 if net.normalization == "batchnorm" else batch)) + counters["per_example_loss_terms"] += ( 2 * config.pert_directions * batch) counters["perturbation_events"] += 1 step += 1 @@ -435,6 +449,10 @@ def parse_args(): parser.add_argument("--weight_decay", type=float, default=5e-4) parser.add_argument("--weight_scale", type=float, default=1.0) parser.add_argument("--residual_scale", type=float) + parser.add_argument("--normalization", choices=("batchnorm", "none"), + default="batchnorm") + parser.add_argument("--bn_momentum", type=float, default=0.1) + parser.add_argument("--bn_eps", type=float, default=1e-5) parser.add_argument("--a_scale", type=float, default=1.0) parser.add_argument("--eta_A", type=float, default=0.01) parser.add_argument("--eta_P", type=float, default=0.01) |