path: root/experiments/conv_local_smoke.py

#!/usr/bin/env python3
"""Prove the convolutional local eligibility matches exact BP when instructed."""
import os
import sys

import torch
import torch.nn.functional as F

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from sdil.conv import (CIFARLocalResNet, CIFARSDILResNet, ConvSDILConfig,
                       conv_local_step, simultaneous_conv_node_perturbation)


def architecture_checks():
    expected = {
        8: (7, 74810),
        20: (19, 268346),
        32: (31, 461882),
        56: (55, 848954),
    }
    for depth, (hidden, parameters) in expected.items():
        net = CIFARLocalResNet(depth=depth)
        assert net.n_hidden == hidden
        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)
        except ValueError:
            pass
        else:
            raise AssertionError(f"invalid depth {depth} was accepted")

    x = torch.arange(2 * 3 * 8 * 8, dtype=torch.float32).reshape(2, 3, 8, 8)
    shortcut = CIFARLocalResNet._option_a_shortcut(x, 6, 2)
    assert tuple(shortcut.shape) == (2, 6, 4, 4)
    assert torch.count_nonzero(shortcut[:, 0]) == 0
    assert torch.count_nonzero(shortcut[:, -2:]) == 0
    assert torch.equal(shortcut[:, 1:4], x[:, :, ::2, ::2])


def exact_local_gradient_check():
    torch.manual_seed(123)
    batch = 3
    x = torch.randn(batch, 3, 32, 32)
    y = torch.tensor([0, 3, 8])
    local = CIFARLocalResNet(depth=8, base_width=4, seed=19)
    bp = CIFARLocalResNet(depth=8, base_width=4, seed=19)

    parameters = local.W + [local.W_out, local.b_out]
    for parameter in parameters:
        parameter.requires_grad_(True)
    forward = local.forward(x, return_cache=True)
    for hidden in forward["hiddens"]:
        hidden.retain_grad()
    loss = F.cross_entropy(forward["logits"], y)
    loss.backward()

    # .backward() differentiated a batch-mean loss.  Multiplying hidden grads
    # by B recovers the per-example convention consumed by the local rule.
    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, 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):
        absolute = (direction + parameter.grad).abs().max()
        scale = parameter.grad.abs().max().clamp_min(1e-12)
        relative_errors.append(float(absolute / scale))
    output_abs = float((out_direction + local.W_out.grad).abs().max())
    bias_abs = float((bias_direction + local.b_out.grad).abs().max())
    assert max(relative_errors) < 3e-5
    assert output_abs < 2e-6 and bias_abs < 2e-6

    for parameter in parameters:
        parameter.requires_grad_(False)
    eta = 0.017
    local.apply_ascent(directions, out_direction, bias_direction, eta)
    bp_loss = bp.bp_step(x, y, eta)
    assert abs(float(loss.detach()) - bp_loss) < 1e-7
    parameter_differences = [
        float((left - right).abs().max())
        for left, right in zip(
            local.W + [local.W_out, local.b_out],
            bp.W + [bp.W_out, bp.b_out])]
    assert max(parameter_differences) < 2e-7
    return {
        "max_relative_local_gradient_error": max(relative_errors),
        "output_absolute_error": output_abs,
        "post_update_parameter_max_error": max(parameter_differences),
    }


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)
    clean = net.forward(x)
    perturbations = [torch.zeros_like(hidden) for hidden in clean["hiddens"]]
    perturbed = net.forward(x, perturbations=perturbations)
    assert torch.equal(clean["logits"], perturbed["logits"])
    perturbations[0] = torch.ones_like(perturbations[0]) * 0.01
    changed = net.forward(x, perturbations=perturbations)
    assert not torch.equal(clean["logits"], changed["logits"])
    try:
        net.forward(x, perturbations=perturbations[:-1])
    except ValueError:
        pass
    else:
        raise AssertionError("short perturbation list was accepted")


def perturbation_estimator_check():
    """Antithetic finite differences equal the simultaneous hidden JVP."""
    torch.manual_seed(3)
    batch = 2
    net = CIFARSDILResNet(
        depth=8, base_width=2, seed=4, dtype=torch.float64)
    x = torch.randn(batch, 3, 32, 32, dtype=torch.float64)
    y = torch.tensor([2, 7])
    parameters = net.W + [net.W_out, net.b_out]
    for parameter in parameters:
        parameter.requires_grad_(True)
    clean = net.forward(x, return_cache=True)
    for hidden in clean["hiddens"]:
        hidden.retain_grad()
    F.cross_entropy(clean["logits"], y).backward()
    generator = torch.Generator(device="cpu").manual_seed(99)
    targets, diagnostics = simultaneous_conv_node_perturbation(
        net, x, y, clean, sigma=1e-6, n_directions=1,
        generator=generator, return_diagnostics=True)
    directions = diagnostics["directions"][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))
    relative = (finite_difference - exact).abs() / exact.abs().clamp_min(1e-12)
    assert float(relative.max()) < 2e-6
    for target, direction in zip(targets, directions):
        expected = -finite_difference[:, None, None, None] * direction
        assert torch.equal(target, expected)
    for parameter in parameters:
        parameter.requires_grad_(False)
    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():
    torch.manual_seed(11)
    net = CIFARSDILResNet(depth=8, base_width=2, seed=6)
    x = torch.randn(8, 3, 32, 32)
    y = torch.arange(8) % 10
    clean = net.forward(x, return_cache=True)
    output_signal = (torch.softmax(clean["logits"], dim=1)
                     - F.one_hot(y, 10))
    prediction, _, _ = net.apical_components(
        output_signal, clean["hiddens"], use_residual=True)
    targets = [torch.randn_like(value) * 0.01 for value in prediction]
    before = sum(float((target - value).square().sum())
                 for target, value in zip(targets, prediction))
    net.calibrate_apical(output_signal, prediction, targets, eta=0.1)
    after_prediction, _, _ = net.apical_components(
        output_signal, clean["hiddens"], use_residual=True)
    after = sum(float((target - value).square().sum())
                for target, value in zip(targets, after_prediction))
    assert after < before

    predictor_net = CIFARSDILResNet(depth=8, base_width=2, seed=8)
    hiddens = [torch.randn(64, *shape) for shape in predictor_net.hidden_shapes]
    initial = predictor_net.predictor_step(hiddens, eta=0.1, nuisance_scale=0.5)
    final = initial
    for _ in range(60):
        final = predictor_net.predictor_step(hiddens, eta=0.1, nuisance_scale=0.5)
    assert final < initial * 1e-3

    weights_before = [weight.clone() for weight in net.W]
    result = conv_local_step(
        net, x[:2], y[:2],
        ConvSDILConfig(
            eta=1e-3, eta_A=1e-3, momentum=0.0, weight_decay=0.0,
            pert_every=1),
        step=0, generator=torch.Generator(device="cpu").manual_seed(7))
    assert result["did_perturb"] and result["calibration"] is not None
    assert torch.isfinite(torch.tensor(list(
        value for key, value in result.items()
        if isinstance(value, float) and key != "predictor_mse"))).all()
    assert any(not torch.equal(before_weight, after_weight)
               for before_weight, after_weight in zip(weights_before, net.W))
    assert all(not parameter.requires_grad
               for parameter in net.W + [net.W_out, net.b_out])
    return {"apical_mse_ratio": after / before,
            "predictor_mse_ratio": final / initial}


def main():
    architecture_checks()
    perturbation_checks()
    report = exact_local_gradient_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")


if __name__ == "__main__":
    main()