File size: 1,461 Bytes
f31f59b | 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 | """Load a paper checkpoint for evaluation; for arm C, round the ternary weights ONCE so inference runs at float speed.
(train.TLinear re-rounds weights on every forward pass, which is right for training and 2x slower for evaluation.)"""
import numpy as np, torch, torch.nn.functional as F
import train
class FrozenTernary(torch.nn.Module):
def __init__(self, lin):
super().__init__()
w = lin.weight.detach()
sw = 1.0 / w.abs().mean().clamp_min(1e-5)
self.register_buffer("wq", ((w * sw).round().clamp(-1, 1) / sw))
self.bias = lin.bias
def forward(self, x):
xf = x.float()
xn = xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + 1e-6)
sx = 127.0 / xn.abs().amax(-1, keepdim=True).clamp_min(1e-5)
xq = (xn * sx).round().clamp(-128, 127) / sx
return F.linear(xq.to(x.dtype), self.wq.to(x.dtype), self.bias)
def load(arm, ckpt, counts, learned_bias=None):
if learned_bias is None:
learned_bias = arm == "C"
m = train.Model(arm, counts, learned_bias=learned_bias).cuda().eval()
st = torch.load(ckpt, map_location="cuda", weights_only=False)
m.load_state_dict(st["model"])
if arm == "C":
for mod in list(m.body.layers.modules()):
for name, child in list(mod.named_children()):
if isinstance(child, train.TLinear):
setattr(mod, name, FrozenTernary(child))
return m, st.get("tag", arm)
|