File size: 1,259 Bytes
47d8ad6 | 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 | import numpy as np, torch, torch.nn as nn, sys
import coremltools as ct
C, S, T = 256, 64, 7
torch.manual_seed(0)
class Stateless(nn.Module):
def __init__(self):
super().__init__()
self.attn = nn.MultiheadAttention(C, 4, batch_first=True)
self.enc = nn.Linear(C, C)
def forward(self, x):
y, _ = self.attn(x, x, x)
return self.enc(y)
class StateNoAttn(nn.Module):
def __init__(self):
super().__init__()
self.register_buffer("bank", torch.zeros(T, S, C))
self.enc = nn.Linear(C, C)
def forward(self, x):
y = self.enc(x) + self.bank.mean(dim=0, keepdim=True)
self.bank.copy_(torch.cat([self.bank[1:], y], 0))
return y
which = sys.argv[1]
m = {"stateless": Stateless, "state_noattn": StateNoAttn}[which]().eval()
m.requires_grad_(False)
x = torch.randn(1, S, C)
with torch.no_grad():
ep = torch.export.export(m, (x,)); ep = ep.run_decompositions({})
ml = ct.convert(ep, minimum_deployment_target=ct.target.iOS18,
compute_units=ct.ComputeUnit.CPU_AND_NE)
st = ml.make_state() if which != "stateless" else None
kw = {"state": st} if st is not None else {}
o = ml.predict({"x": x.numpy()}, **kw)
print(which, "on CPU_AND_NE: predict OK")
|