scripts/check_torch.py

"""Cross-check: our left-to-right FP simulator vs PyTorch autograd on random networks."""
import numpy as np, torch, json, sys
from fpnet import Act, FPNet

TORCH_ACT = {"relu": torch.relu, "elu": torch.nn.functional.elu,
             "gelu": lambda t: torch.nn.functional.gelu(t, approximate="none"),
             "swish": torch.nn.functional.silu, "sigmoid": torch.sigmoid, "tanh": torch.tanh}


def main(seed=0):
    rng = np.random.default_rng(seed)
    rows = []
    for name in TORCH_ACT:
        for dt, tdt in ((np.float32, torch.float32), (np.float64, torch.float64)):
            act = Act(name, dt)
            dims = [3, 5, 4, 1]
            As = [rng.normal(size=(dims[i + 1], dims[i])).astype(dt) for i in range(3)]
            bs = [rng.normal(size=dims[i + 1]).astype(dt) for i in range(3)]
            net = FPNet(As, bs, act)
            x = rng.normal(size=3).astype(dt)
            g0 = dt(1.7)
            v, g = net.value_and_grad(x, g0)
            tx = torch.tensor(x, dtype=tdt, requires_grad=True)
            h = tx
            for i in range(3):
                h = torch.tensor(As[i], dtype=tdt) @ h + torch.tensor(bs[i], dtype=tdt)
                if i < 2:
                    h = TORCH_ACT[name](h)
            tg = torch.autograd.grad(h * float(g0), tx)[0]
            rows.append(dict(act=name, dtype=dt.__name__,
                             val_abs_err=float(abs(float(v[0]) - float(h.item()))),
                             grad_max_abs_err=float(np.max(np.abs(g - tg.numpy()))),
                             grad_max_rel_err=float(np.max(np.abs(
                                 (g - tg.numpy()) / np.maximum(np.abs(tg.numpy()), 1e-30))))))
    return rows


if __name__ == "__main__":
    rows = main()
    for r in rows:
        print(r)
    import os
    os.makedirs("../outputs", exist_ok=True)
    with open("../outputs/torch_crosscheck.json", "w") as f:
        json.dump(rows, f, indent=1)