diff options
Diffstat (limited to 'experiments/conv_local_smoke.py')
| -rw-r--r-- | experiments/conv_local_smoke.py | 129 |
1 files changed, 125 insertions, 4 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") |