File size: 2,028 Bytes
309b968 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 | """Shared helpers for verified neural units.
Same discipline as the neural-aarch64 / neural-photonic / neural-ddr units:
a small MLP is trained until it is *bit-identical to a golden reference over its
entire finite input domain* (N/N verification).
"""
from __future__ import annotations
import torch
import torch.nn as nn
def bits_of(v: int, n: int) -> torch.Tensor:
"""LSB-first bit vector of length n."""
return torch.tensor([(v >> i) & 1 for i in range(n)], dtype=torch.float32)
def int_of(bits: torch.Tensor) -> int:
"""Inverse of bits_of: LSB-first bit vector -> int."""
v = 0
for i, b in enumerate(bits.tolist()):
if b >= 0.5:
v |= (1 << i)
return v
def pm(bits: torch.Tensor) -> torch.Tensor:
"""Map {0,1} bits to {-1,+1} for a friendlier input scale."""
return bits * 2.0 - 1.0
def mlp(inp: int, out: int, h: int = 128, layers: int = 3) -> nn.Sequential:
mods: list[nn.Module] = [nn.Linear(inp, h), nn.ReLU()]
for _ in range(layers - 1):
mods += [nn.Linear(h, h), nn.ReLU()]
mods += [nn.Linear(h, out)]
return nn.Sequential(*mods)
@torch.no_grad()
def verify(net: nn.Module, X: torch.Tensor, Ybits: torch.Tensor) -> tuple[int, int]:
"""Return (n_correct, n_total) over the full enumerated domain."""
net.eval()
pred = (net(X) > 0).float()
ok = (pred == Ybits).all(dim=1).sum().item()
return int(ok), X.shape[0]
def train(net: nn.Module, X: torch.Tensor, Ybits: torch.Tensor,
steps: int = 4000, lr: float = 2e-3, tag: str = "") -> nn.Module:
opt = torch.optim.Adam(net.parameters(), lr=lr)
lossf = nn.BCEWithLogitsLoss()
net.train()
for s in range(steps):
opt.zero_grad()
loss = lossf(net(X), Ybits)
loss.backward()
opt.step()
if tag and (s % 1000 == 0 or s == steps - 1):
ok, tot = verify(net, X, Ybits)
net.train()
print(f" [{tag}] step {s:5d} loss {loss.item():.5f} verify {ok}/{tot}")
return net
|