repro_bundle/code/efficiency.py

"""Claim 6 (mechanism-level) + Claim 3 (param count) efficiency microbenchmark.

The paper's headline efficiency numbers (45% faster, 18% less memory than LGD on
3DPW, Table 2) come from the *full* Human-Mesh-Recovery pipeline, which we cannot
run here (needs AMASS + SMPL + 3DPW + the LGD framework + 400K meta-iters). What
we CAN test is the mechanism the claim rests on: L-SR1's update rule is three
small *element-wise* MLPs (d_hidden=128) plus a low-rank buffer, i.e. its per-step
cost is independent of the number of optimized parameters, whereas an LGD-style
learned updater regresses the update from the whole state and grows with it.

We therefore compare, per inner optimization step, the runtime and peak memory of:
  * the L-SR1 update module (encoder + vector-gen + lr-gen, element-wise), and
  * a plausible LGD-style updater (a fully-connected regressor over the full
    state), sized to the same order of parameters.

This is a *proxy* for the direction of Claim 6, NOT a reproduction of the exact
3DPW numbers. Everything is labelled accordingly in the logbook.
"""
from __future__ import annotations

import argparse
import json
import time

import torch
import torch.nn as nn

from lsr1 import LSR1Config, LSR1Optimizer


class LGDStyleUpdater(nn.Module):
    """LGD (Song et al. 2020) style learned update: a fully-connected regressor
    that maps the concatenated state (params, gradient) of the WHOLE problem to a
    parameter update. Width does not shrink with dimension -> cost grows with the
    number of optimized parameters, unlike L-SR1's element-wise modules."""

    def __init__(self, dim, hidden=1024, layers=3):
        super().__init__()
        blocks = [nn.Linear(2 * dim, hidden), nn.ReLU()]
        for _ in range(layers - 1):
            blocks += [nn.Linear(hidden, hidden), nn.ReLU()]
        blocks += [nn.Linear(hidden, dim)]
        self.net = nn.Sequential(*blocks)

    def forward(self, x, g):
        return self.net(torch.cat([x, g], dim=-1))


def lsr1_step(opt, x, buffer):
    """One L-SR1 update-module forward on state x [B, N] (element-wise):
    3 element-wise MLP forwards (encoder, vector-gen, lr-gen) + low-rank buffer
    sum over the L stored vectors."""
    g = torch.randn_like(x)
    p = torch.zeros_like(x); d_prev = torch.zeros_like(x); q = torch.zeros_like(x)
    v, alpha, d, buffer = opt.step_direction(x, p, d_prev, g, q, buffer)
    return x - alpha * d, buffer


def bench(fn, warmup=5, iters=30, device="cpu"):
    cuda = device.startswith("cuda")
    for _ in range(warmup):
        fn()
    if cuda:
        torch.cuda.synchronize(); torch.cuda.reset_peak_memory_stats()
    t0 = time.time()
    for _ in range(iters):
        fn()
    if cuda:
        torch.cuda.synchronize()
    dt = (time.time() - t0) / iters * 1000.0        # ms/iter
    mem = (torch.cuda.max_memory_allocated() / 2**30) if cuda else float("nan")
    return dt, mem


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--device", default="cpu")
    ap.add_argument("--batch", type=int, default=256)   # paper Table 2 batch
    ap.add_argument("--dim", type=int, default=256,
                    help="number of optimized parameters per sample")
    ap.add_argument("--buffer", type=int, default=4)    # HMR buffer L=4
    ap.add_argument("--out", default="outputs/efficiency.json")
    args = ap.parse_args()
    dev = args.device
    torch.manual_seed(0)

    cfg = LSR1Config(buffer_L=args.buffer, d_hidden=128)
    lsr1 = LSR1Optimizer(cfg).to(dev).eval()
    lgd = LGDStyleUpdater(args.dim, hidden=1024, layers=3).to(dev).eval()

    x = torch.randn(args.batch, args.dim, device=dev)
    # Pre-fill the L-vector buffer ONCE so we time a single steady-state inner
    # step (not the warm-up), matching the paper's per-inner-iteration Table 2.
    steady_buf = [torch.randn(args.batch, args.dim, device=dev)
                  for _ in range(args.buffer)]

    def run_lsr1():
        _, _ = lsr1_step(lsr1, x, list(steady_buf))
        return _

    def run_lgd():
        g = torch.randn_like(x)
        return lgd(x, g)

    with torch.no_grad():
        t_lsr1, m_lsr1 = bench(run_lsr1, device=dev)
        t_lgd, m_lgd = bench(run_lgd, device=dev)

    p_lsr1 = sum(p.numel() for p in lsr1.parameters())
    p_lgd = sum(p.numel() for p in lgd.parameters())
    res = dict(
        device=dev, batch=args.batch, dim=args.dim, buffer=args.buffer,
        lsr1_ms=t_lsr1, lgd_ms=t_lgd,
        lsr1_mem_gib=m_lsr1, lgd_mem_gib=m_lgd,
        speedup_pct=100 * (1 - t_lsr1 / t_lgd),
        mem_reduction_pct=(100 * (1 - m_lsr1 / m_lgd)
                           if m_lgd == m_lgd else None),
        lsr1_update_params=p_lsr1, lgd_update_params=p_lgd,
        paper_reference=dict(lgd_ms=166, lsr1_ms=91, lgd_mem_gib=17.81,
                             lsr1_mem_gib=14.60, speedup_pct=45,
                             mem_reduction_pct=18,
                             lgd_full_params=17.4e6, lsr1_full_params=10.4e6),
    )
    with open(args.out, "w") as f:
        json.dump(res, f, indent=2)
    print(json.dumps({k: res[k] for k in
                      ("lsr1_ms", "lgd_ms", "speedup_pct", "lsr1_mem_gib",
                       "lgd_mem_gib", "mem_reduction_pct",
                       "lsr1_update_params", "lgd_update_params")}, indent=2))


if __name__ == "__main__":
    main()