scripts/bench_forward_cache.py - master - Mihaham/CNN-NEAT

scripts/bench_forward_cache.py
463 строки14 KB
MihahamYT
feat(scripts): OVA holdout, hash studies, and cache benches
04 авг 2026, 20:36
04 авг 2026, 20:36e256d0f
#!/usr/bin/env python3
"""Load-test ForwardActivationCache vs ScreenCacheStore (CPU, synthetic NEAT pop)."""

from __future__ import annotations

import argparse
import json
import random
import time
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Dict, List, Optional, Tuple

import torch

from cnn_neat.forward_cache import ForwardActivationCache, data_id_from_tensor
from cnn_neat.lineage import GenomeLineage
from cnn_neat.mutations import ConvolutionalNetworkMutator, MutationType
from cnn_neat.network import ConvolutionalNeuralNetwork
from cnn_neat.screen_cache import ScreenCacheStore


GenomeRec = Tuple[ConvolutionalNeuralNetwork, GenomeLineage, Tuple[int, int]]


@dataclass
class BenchResult:
    name: str
    n_genomes: int
    wall_s: float
    bytes_used: int
    peak_bytes: int
    hits: int
    misses: int
    incremental_hits: int
    exact_hits: int
    evictions: int
    entries: int
    shared_nodes: int
    dup_bytes_saved: int
    max_abs_err: float
    hit_rate: float


def _seed_genome(h: int = 8, w: int = 8, out: int = 4) -> ConvolutionalNeuralNetwork:
    g = ConvolutionalNeuralNetwork(
        input_image_size=(h, w),
        output_image_size=(out, out),
        num_input_views=1,
        minimal_connections=True,
        activation=None,
        edge_weights=False,
    )
    # Grow a small shared stem
    for _ in range(2):
        keys = [k for k, c in g.connections.items() if c.enabled]
        if not keys:
            break
        hid = g.add_node_split(keys[0])
        if hid == -1:
            break
    return g


def _build_population(
    n: int,
    *,
    seed: int,
    n_elites: int,
) -> Tuple[List[GenomeRec], List[int]]:
    """
    Return population records and elite indices.

    Elites share topology but have distinct weights (separate ScreenCache
    entries). Offspring are mostly clones/splits from elites so the hash
    cache can reuse node activations across individuals.
    """
    rng = random.Random(seed)
    torch.manual_seed(seed)
    stem = _seed_genome()
    elites: List[ConvolutionalNeuralNetwork] = []
    for i in range(max(1, n_elites)):
        e = stem.clone()
        # Distinct weight fingerprint per elite → separate ScreenCacheStore keys
        mut = ConvolutionalNetworkMutator(e)
        elites.append(
            mut.apply(MutationType.MUTATE_WEIGHTS, rate=1.0, scale=0.05 + 0.01 * i)
        )

    pop: List[GenomeRec] = []
    elite_indices: List[int] = []
    for i, e in enumerate(elites):
        gid = (0, i)
        pop.append(
            (
                e,
                GenomeLineage(origin="elite", mutation_type=None),
                gid,
            )
        )
        elite_indices.append(i)

    # Bias toward clone / split so shared prefixes dominate
    ops = [
        MutationType.CLONE,
        MutationType.CLONE,
        MutationType.CLONE,
        MutationType.SPLIT,
        MutationType.WITHOUT_SPLIT,
        MutationType.MUTATE_WEIGHTS,
    ]
    while len(pop) < n:
        parent_idx = elite_indices[rng.randrange(len(elite_indices))]
        parent, _, parent_gid = pop[parent_idx]
        op = ops[rng.randrange(len(ops))]
        child = ConvolutionalNetworkMutator(parent).apply(op)
        idx = len(pop)
        lineage = GenomeLineage(
            origin="offspring",
            parent_generation=parent_gid[0],
            parent_index=parent_gid[1],
            mutation_type=op.value,
        )
        pop.append((child, lineage, (1, idx)))
    return pop, elite_indices


def _max_err(a: torch.Tensor, b: torch.Tensor) -> float:
    return float((a.float() - b.float()).abs().max().item())


@torch.inference_mode()
def _run_screen_store(
    pop: List[GenomeRec],
    elite_indices: List[int],
    images: torch.Tensor,
    max_bytes: int,
) -> Tuple[BenchResult, Dict[Tuple[int, int], torch.Tensor]]:
    store = ScreenCacheStore(max_entries=10_000, max_bytes=max_bytes)
    outputs: Dict[Tuple[int, int], torch.Tensor] = {}
    max_err = 0.0
    t0 = time.perf_counter()

    # Refresh elites (full capture)
    for ei in elite_indices:
        genome, _, gid = pop[ei]
        genome.eval()
        out, node_acts = genome.forward_execute_with_capture(images)
        store.put(
            genome,
            generation=gid[0],
            population_index=gid[1],
            node_activations=node_acts,
            output_features=out,
        )
        outputs[gid] = out
        base = genome(images)
        max_err = max(max_err, _max_err(out, base))

    exact_hits = 0
    incr_hits = 0
    misses = 0
    for i, (genome, lineage, gid) in enumerate(pop):
        if i in elite_indices:
            continue
        genome.eval()
        exact = store.get_exact(genome)
        if exact is not None:
            exact_hits += 1
            outputs[gid] = exact.output_features.float()
            continue
        parent_entry, start_step = store.incremental_start_step(genome, lineage)
        plan = genome.execution_plan
        if (
            parent_entry is not None
            and plan is not None
            and start_step < len(plan.merge_steps)
        ):
            out = genome.forward_execute_incremental(
                images,
                parent_node_activations={
                    k: v.float() for k, v in parent_entry.node_activations.items()
                },
                start_merge_step=start_step,
            )
            incr_hits += 1
            outputs[gid] = out
            base = genome(images)
            max_err = max(max_err, _max_err(out, base))
            continue
        store.note_miss()
        misses += 1
        out = genome(images)
        outputs[gid] = out

    wall = time.perf_counter() - t0
    hits = store.stats.exact_hits + store.stats.incremental_hits
    total = hits + store.stats.misses
    return (
        BenchResult(
            name="ScreenCacheStore",
            n_genomes=len(pop),
            wall_s=wall,
            bytes_used=int(store.stats.bytes_used),
            peak_bytes=int(store.stats.bytes_used),
            hits=int(hits),
            misses=int(store.stats.misses),
            incremental_hits=int(store.stats.incremental_hits),
            exact_hits=int(store.stats.exact_hits),
            evictions=int(store.stats.evictions),
            entries=len(store),
            shared_nodes=0,
            dup_bytes_saved=0,
            max_abs_err=max_err,
            hit_rate=(float(hits) / float(total)) if total else 0.0,
        ),
        outputs,
    )


@torch.inference_mode()
def _run_hash_cache(
    pop: List[GenomeRec],
    elite_indices: List[int],
    images: torch.Tensor,
    max_bytes: int,
    *,
    mode: str = "legacy",
) -> Tuple[BenchResult, Dict[Tuple[int, int], torch.Tensor]]:
    """
    mode:
      legacy — every genome stores activations during the eval pass.
      elite_commit — two-wave protocol matching evolution_loop:
          wave0: precompute + eval(no store) + commit elites
          wave1: offspring from elites, precompute + eval(no store) using
                 the elite pool (hits), then re-commit elites
    """
    data_id = data_id_from_tensor(images)
    cache = ForwardActivationCache(max_bytes=max_bytes, data_id=data_id)
    watch = {pop[i][2] for i in elite_indices}
    cache.set_watchlist(watch)
    outputs: Dict[Tuple[int, int], torch.Tensor] = {}
    max_err = 0.0
    t0 = time.perf_counter()

    if mode == "elite_commit":
        # --- generation 0 ---
        cache.precompute_population([(g, gid) for g, _l, gid in pop])
        for i, (genome, _lin, gid) in enumerate(pop):
            genome.eval()
            out = cache.execute(
                genome, images, gid, pin=False, store_activations=False
            )
            outputs[gid] = out
            base = genome(images)
            max_err = max(max_err, _max_err(out, base))
        for ei in elite_indices:
            genome, _lin, gid = pop[ei]
            genome.eval()
            cache.execute(genome, images, gid, pin=True, store_activations=True)

        # --- generation 1: children of elites (reuse pool) ---
        gen1: List[GenomeRec] = []
        for j, ei in enumerate(elite_indices):
            parent, _plin, pgid = pop[ei]
            child = ConvolutionalNetworkMutator(parent).apply(MutationType.CLONE)
            gid = (2, j)
            gen1.append(
                (
                    child,
                    GenomeLineage(
                        origin="offspring",
                        parent_generation=pgid[0],
                        parent_index=pgid[1],
                        mutation_type="clone",
                    ),
                    gid,
                )
            )
            mut = ConvolutionalNetworkMutator(parent).apply(MutationType.SPLIT)
            gid_s = (2, len(elite_indices) + j)
            gen1.append(
                (
                    mut,
                    GenomeLineage(
                        origin="offspring",
                        parent_generation=pgid[0],
                        parent_index=pgid[1],
                        mutation_type="split",
                    ),
                    gid_s,
                )
            )
        cache.precompute_population([(g, gid) for g, _l, gid in gen1])
        for genome, _lin, gid in gen1:
            genome.eval()
            out = cache.execute(
                genome, images, gid, pin=False, store_activations=False
            )
            outputs[gid] = out
            base = genome(images)
            max_err = max(max_err, _max_err(out, base))
        # Re-commit a subset as next breeding pool
        for genome, _lin, gid in gen1[: max(1, len(elite_indices))]:
            genome.eval()
            cache.execute(genome, images, gid, pin=True, store_activations=True)
    else:
        for i, (genome, _lin, gid) in enumerate(pop):
            genome.eval()
            pin = i in elite_indices
            out = cache.execute(
                genome, images, gid, pin=pin, store_activations=True
            )
            outputs[gid] = out
            base = genome(images)
            max_err = max(max_err, _max_err(out, base))

    wall = time.perf_counter() - t0
    st = cache.stats
    name = (
        "ForwardHashCache_elite_commit"
        if mode == "elite_commit"
        else "ForwardActivationCache"
    )
    return (
        BenchResult(
            name=name,
            n_genomes=len(pop),
            wall_s=wall,
            bytes_used=int(st.bytes_used),
            peak_bytes=int(st.peak_bytes),
            hits=int(st.hits),
            misses=int(st.misses),
            incremental_hits=0,
            exact_hits=int(st.hits),
            evictions=int(st.evictions),
            entries=int(st.entries),
            shared_nodes=int(st.shared_nodes),
            dup_bytes_saved=int(st.dup_bytes_saved),
            max_abs_err=max_err,
            hit_rate=float(st.hit_rate()),
        ),
        outputs,
    )


def _print_table(results: List[BenchResult]) -> None:
    cols = [
        ("name", 28),
        ("wall_s", 10),
        ("bytes", 12),
        ("peak", 12),
        ("hits", 8),
        ("miss", 8),
        ("evict", 8),
        ("entries", 8),
        ("shared", 8),
        ("hit%", 8),
        ("max_err", 10),
    ]
    header = " ".join(f"{h:>{w}}" for h, w in cols)
    print(header)
    print("-" * len(header))
    for r in results:
        row = [
            f"{r.name:>28}",
            f"{r.wall_s:>10.3f}",
            f"{r.bytes_used:>12}",
            f"{r.peak_bytes:>12}",
            f"{r.hits:>8}",
            f"{r.misses:>8}",
            f"{r.evictions:>8}",
            f"{r.entries:>8}",
            f"{r.shared_nodes:>8}",
            f"{100.0 * r.hit_rate:>7.1f}%",
            f"{r.max_abs_err:>10.4g}",
        ]
        print(" ".join(row))


def run_bench(
    *,
    n: int = 200,
    n_elites: int = 10,
    batch: int = 8,
    max_bytes: int = 32_000_000,
    seed: int = 0,
    out_json: Optional[Path] = None,
) -> List[BenchResult]:
    pop, elite_indices = _build_population(n, seed=seed, n_elites=n_elites)
    torch.manual_seed(seed + 1)
    images = torch.randn(batch, 1, 3, 8, 8)

    screen_res, _ = _run_screen_store(pop, elite_indices, images, max_bytes)
    hash_legacy, _ = _run_hash_cache(
        pop, elite_indices, images, max_bytes, mode="legacy"
    )
    hash_elite, _ = _run_hash_cache(
        pop, elite_indices, images, max_bytes, mode="elite_commit"
    )
    results = [screen_res, hash_legacy, hash_elite]
    _print_table(results)
    if hash_legacy.wall_s > 0:
        speedup = hash_legacy.wall_s / max(hash_elite.wall_s, 1e-9)
        print(
            f"\nelite_commit vs legacy store-all: "
            f"{speedup:.2f}x "
            f"(legacy={hash_legacy.wall_s:.3f}s, "
            f"elite={hash_elite.wall_s:.3f}s, "
            f"bytes {hash_legacy.bytes_used} -> {hash_elite.bytes_used})"
        )

    if out_json is not None:
        payload = {
            "config": {
                "n": n,
                "n_elites": n_elites,
                "batch": batch,
                "max_bytes": max_bytes,
                "seed": seed,
            },
            "results": [asdict(r) for r in results],
            "compare": {
                "elite_vs_legacy_speedup": (
                    hash_legacy.wall_s / max(hash_elite.wall_s, 1e-9)
                ),
                "legacy_wall_s": hash_legacy.wall_s,
                "elite_wall_s": hash_elite.wall_s,
                "legacy_bytes": hash_legacy.bytes_used,
                "elite_bytes": hash_elite.bytes_used,
            },
        }
        out_json.parent.mkdir(parents=True, exist_ok=True)
        out_json.write_text(json.dumps(payload, indent=2), encoding="utf-8")
        print(f"Wrote {out_json}")
    return results


def main() -> None:
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--n", type=int, default=200)
    p.add_argument("--n-elites", type=int, default=10)
    p.add_argument("--batch", type=int, default=8)
    p.add_argument("--max-bytes", type=int, default=32_000_000)
    p.add_argument("--seed", type=int, default=0)
    p.add_argument(
        "--out",
        type=Path,
        default=Path("scripts/bench_forward_cache_summary.json"),
    )
    args = p.parse_args()
    run_bench(
        n=args.n,
        n_elites=args.n_elites,
        batch=args.batch,
        max_bytes=args.max_bytes,
        seed=args.seed,
        out_json=args.out,
    )


if __name__ == "__main__":
    main()