path: root/experiments/resnet_crossover_smoke.py

#!/usr/bin/env python3
"""Deterministic equation audits for matched ResNet crossover adapters."""
import json
import os
import sys

import torch

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from sdil.conv import CIFARLocalResNet
from sdil.conv_crossover import (
    CIFARDualPropResNet,
    CIFARForwardForwardResNet,
    CIFARPEPITAResNet,
)


def relative_error(actual, expected):
    numerator = torch.linalg.vector_norm(actual - expected)
    denominator = torch.linalg.vector_norm(expected).clamp_min(1e-30)
    return float(numerator / denominator)


def audit_dualprop():
    common = dict(
        depth=8, base_width=2, seed=71, normalization="batchnorm",
        residual_scale=1.0, dtype=torch.float64)
    reference = CIFARLocalResNet(**common)
    net = CIFARDualPropResNet(
        **common, alpha=0.0, dp_beta=0.1, inference_passes=2)
    generator = torch.Generator().manual_seed(72)
    image = torch.randn(3, 3, 32, 32, generator=generator,
                        dtype=torch.float64)
    labels = torch.tensor([0, 3, 7])
    one_hot = torch.nn.functional.one_hot(labels, 10).to(torch.float64)
    with torch.no_grad():
        reference_output = reference.forward(
            image, training=True, update_stats=False)["logits"]
        net_output = net.forward(
            image, training=True, update_stats=False)["logits"]
        clean = net.forward(
            image, return_cache=True, training=True, update_stats=False)
        plus, minus = net.infer_dual_states(image, one_hot, clean)
    forward_error = float(torch.max(torch.abs(
        reference_output - net_output)))

    parameters = (
        net.W + net.gamma + net.beta + [net.W_out, net.b_out])
    for parameter in parameters:
        parameter.requires_grad_(True)
    alpha = net.dp_alpha
    beta = net.dp_beta
    states = [
        (alpha * positive + (1.0 - alpha) * negative).detach()
        for positive, negative in zip(plus[:-1], minus[:-1])]
    deltas = [
        ((positive - negative) / beta).detach()
        for positive, negative in zip(plus, minus)]
    objective = image.new_zeros(())
    for index in range(net.n_hidden):
        prediction, _ = net._node_prediction(index, states, image)
        objective -= torch.sum(deltas[index] * prediction) / image.shape[0]
    features = states[-1].mean(dim=(2, 3))
    output_prediction = features @ net.W_out.t() + net.b_out
    objective -= torch.sum(deltas[-1] * output_prediction) / image.shape[0]
    gradients = torch.autograd.grad(objective, parameters)
    (directions, gamma_directions, beta_directions,
     output_weight, output_bias) = net.dualprop_ascent_directions(
         image, plus, minus)
    actual = (
        directions + gamma_directions + beta_directions
        + [output_weight, output_bias])
    errors = [
        relative_error(direction, -gradient)
        for direction, gradient in zip(actual, gradients)]
    for parameter in parameters:
        parameter.requires_grad_(False)
    if forward_error >= 1e-12 or max(errors) >= 2e-12:
        raise AssertionError({
            "forward_error": forward_error,
            "direction_errors": errors,
        })
    return {
        "matched_forward_max_absolute_error": forward_error,
        "contrastive_direction_max_relative_error": max(errors),
        "num_audited_parameter_tensors": len(errors),
        "uses_symmetric_forward_edge_transposes": True,
        "uses_reverse_task_loss_graph": False,
    }


def audit_pepita():
    common = dict(
        depth=8, base_width=2, seed=81, normalization="batchnorm",
        residual_scale=1.0, dtype=torch.float64)
    reference = CIFARLocalResNet(**common)
    net = CIFARPEPITAResNet(
        **common, projection_scale=0.05, projection_seed=82)
    generator = torch.Generator().manual_seed(83)
    image = torch.randn(3, 3, 32, 32, generator=generator,
                        dtype=torch.float64)
    labels = torch.tensor([1, 4, 8])
    one_hot = torch.nn.functional.one_hot(labels, 10).to(torch.float64)
    with torch.no_grad():
        reference_output = reference.forward(
            image, training=True, update_stats=False)["logits"]
        clean = net.forward(
            image, return_cache=True, training=True, update_stats=False)
        clean_error = torch.softmax(clean["logits"], dim=1) - one_hot
        input_error = torch.einsum(
            "bc,cijk->bijk", clean_error, net.input_feedback)
        modulated = net.forward(
            image + input_error, return_cache=True,
            training=True, update_stats=False)
        modulated_error = (
            torch.softmax(modulated["logits"], dim=1) - one_hot)
    forward_error = float(torch.max(torch.abs(
        reference_output - clean["logits"])))

    parameters = (
        net.W + net.gamma + net.beta + [net.W_out, net.b_out])
    for parameter in parameters:
        parameter.requires_grad_(True)
    objective = image.new_zeros(())
    batch = image.shape[0]
    for index, (clean_hidden, modulated_hidden, cache, spec) in enumerate(zip(
            clean["hiddens"], modulated["hiddens"], modulated["caches"],
            net.layer_specs)):
        field = (clean_hidden - modulated_hidden).detach()
        convolution = torch.nn.functional.conv2d(
            cache["pre"].detach(), net.W[index],
            stride=spec.stride, padding=spec.padding)
        normalized, _ = net._normalize(
            index, convolution, training=True, update_stats=False)
        prediction = spec.branch_scale * normalized
        spatial = prediction.shape[2] * prediction.shape[3]
        objective += torch.sum(field * prediction) / (batch * spatial)
    output_prediction = (
        modulated["features"].detach() @ net.W_out.t() + net.b_out)
    objective += torch.sum(
        modulated_error.detach() * output_prediction) / batch
    gradients = torch.autograd.grad(objective, parameters)
    (directions, gamma_directions, beta_directions,
     output_weight, output_bias) = net.pepita_ascent_directions(
         clean, modulated, modulated_error)
    actual = (
        directions + gamma_directions + beta_directions
        + [output_weight, output_bias])
    errors = [
        relative_error(direction, -gradient)
        for direction, gradient in zip(actual, gradients)]
    for parameter in parameters:
        parameter.requires_grad_(False)
    projection_limit = (6.0 / (3 * 32 * 32)) ** 0.5 * 0.05
    observed_limit = float(torch.max(torch.abs(net.input_feedback)))
    if (forward_error >= 1e-12 or max(errors) >= 2e-12
            or observed_limit > projection_limit):
        raise AssertionError({
            "forward_error": forward_error,
            "direction_errors": errors,
            "projection_limit": projection_limit,
            "observed_limit": observed_limit,
        })
    return {
        "matched_forward_max_absolute_error": forward_error,
        "local_equation_max_relative_error": max(errors),
        "input_projection_shape": list(net.input_feedback.shape),
        "input_projection_limit": projection_limit,
        "observed_input_projection_max": observed_limit,
        "uses_reverse_task_loss_graph": False,
    }


def audit_forward_forward():
    common = dict(
        depth=8, base_width=2, seed=91, normalization="batchnorm",
        residual_scale=1.0, dtype=torch.float64)
    reference = CIFARLocalResNet(**common)
    net = CIFARForwardForwardResNet(
        **common, threshold=2.0, learning_rate=0.03,
        score_from_layer=1)
    if net.n_forward_parameters != reference.n_forward_parameters:
        raise AssertionError("Forward-Forward changed forward parameter count")
    generator = torch.Generator().manual_seed(92)
    image = torch.randn(4, 3, 32, 32, generator=generator,
                        dtype=torch.float64)
    labels = torch.tensor([1, 2, 6, 9])
    negative_labels = (labels + 3) % 10
    target = 1
    with torch.no_grad():
        positive = net.ff_forward(
            net.ff_overlay(image, labels),
            training=True, update_stats=False)
        negative = net.ff_forward(
            net.ff_overlay(image, negative_labels),
            training=True, update_stats=False)
    positive_output, negative_output = net._ff_local_outputs(
        target, positive, negative)
    axes = tuple(range(1, positive_output.ndim))
    positive_goodness = positive_output.square().mean(dim=axes)
    negative_goodness = negative_output.square().mean(dim=axes)
    explicit_loss = (
        torch.nn.functional.softplus(-positive_goodness + net.ff_threshold)
        + torch.nn.functional.softplus(
            negative_goodness - net.ff_threshold)
    ).mean()
    parameters = (
        net.W + net.gamma + net.beta + [net.W_out, net.b_out])
    before = [parameter.detach().clone() for parameter in parameters]
    metrics = net.ff_train_layer(
        target, image, labels, learning_rate=0.03,
        negative_labels=negative_labels)
    changed = [
        not torch.equal(old, parameter.detach())
        for old, parameter in zip(before, parameters)]
    target_indices = {target}
    if net.normalization == "batchnorm":
        target_indices.update({
            len(net.W) + target,
            len(net.W) + len(net.gamma) + target,
        })
    non_target_changes = [
        index for index, value in enumerate(changed)
        if value and index not in target_indices]
    non_target_gradients = [
        index for index, parameter in enumerate(parameters)
        if index not in target_indices and parameter.grad is not None]
    scores = net.ff_candidate_scores(image)
    loss_error = abs(metrics["loss"] - float(explicit_loss.detach()))
    if (loss_error >= 1e-12 or non_target_changes
            or non_target_gradients or not any(
                changed[index] for index in target_indices)
            or scores.shape != (4, 10)
            or not torch.isfinite(scores).all()):
        raise AssertionError({
            "loss_error": loss_error,
            "non_target_changes": non_target_changes,
            "non_target_gradients": non_target_gradients,
            "changed": changed,
            "score_shape": list(scores.shape),
        })
    return {
        "local_objective_absolute_error": loss_error,
        "non_target_parameter_changes": len(non_target_changes),
        "non_target_gradients": len(non_target_gradients),
        "candidate_score_shape": list(scores.shape),
        "matched_forward_parameter_count": net.n_forward_parameters,
        "uses_only_target_layer_autograd": True,
    }


def main():
    print(json.dumps({
        "dualprop": audit_dualprop(),
        "forward_forward": audit_forward_forward(),
        "pepita": audit_pepita(),
    },
                     indent=2, sort_keys=True))


if __name__ == "__main__":
    main()