File size: 3,600 Bytes
8a5e2d4 | 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 65 66 67 68 69 70 71 72 73 74 75 76 77 78 | """Export a package folder (config.json, backbone/, heads.safetensors) to ONNX for the in-browser demo.
python onnx_export.py DIR # writes DIR/onnx/model_fp32.onnx and DIR/onnx/model_int8.onnx
The graph takes input_ids and attention_mask (int64, dynamic batch and sequence length) and returns `probs`
[batch, n_heads], the sigmoid of each head's logit, in the order of config.json "heads". The int8 file is dynamic
quantisation of the weights with per-channel scales: the embedding table (19.3 of the 32 million parameters) and the
attention and MLP up-projection matrices. The MLP down-projection stays fp32: quantising it cost up to 0.02 AUROC
on the delegate head, the other layers cost under 0.001 (see PARITY.md).
"""
from __future__ import annotations
import argparse
import json
import sys
from pathlib import Path
import torch
sys.path.insert(0, str(Path(__file__).resolve().parent))
from cue import CueModel # noqa: E402
class ProbWrapper(torch.nn.Module):
def __init__(self, m: CueModel):
super().__init__()
self.backbone = m.backbone
self.w = torch.nn.Parameter(torch.cat([m.heads[q].weight for q in m.head_names], 0).detach().clone(), requires_grad=False)
self.b = torch.nn.Parameter(torch.cat([m.heads[q].bias for q in m.head_names], 0).detach().clone(), requires_grad=False)
def forward(self, input_ids, attention_mask):
b = self.backbone
h = b.model(input_ids=input_ids, attention_mask=attention_mask)[0]
mk = attention_mask.unsqueeze(-1).to(h.dtype)
h = (h * mk).sum(1) / mk.sum(1) # mean pooling (config classifier_pooling = "mean")
v = b.head(h)
return torch.sigmoid(v @ self.w.T + self.b)
def export(folder: Path, opset: int = 17):
from transformers import AutoModelForSequenceClassification
from onnxruntime.quantization import QuantType, quantize_dynamic
folder = Path(folder)
cfg = json.loads((folder / "config.json").read_text())
m = CueModel.from_pretrained(folder)
if m.backbone.config.classifier_pooling != "mean":
sys.exit("only mean pooling is exported")
# eager attention exports cleanly; the weights are identical
m.backbone = AutoModelForSequenceClassification.from_pretrained(folder / cfg["backbone_dir"], attn_implementation="eager").eval()
w = ProbWrapper(m).eval()
out = folder / "onnx"
out.mkdir(exist_ok=True)
ids = torch.randint(5, 1000, (2, 24))
mask = torch.ones_like(ids)
mask[1, 18:] = 0
fp32 = out / "model_fp32.onnx"
torch.onnx.export(w, (ids, mask), str(fp32), input_names=["input_ids", "attention_mask"], output_names=["probs"],
dynamic_axes={"input_ids": {0: "batch", 1: "seq"}, "attention_mask": {0: "batch", 1: "seq"},
"probs": {0: "batch"}}, opset_version=opset, dynamo=False)
int8 = out / "model_int8.onnx"
import onnx
names = [n.name for n in onnx.load(str(fp32)).graph.node if n.op_type == "MatMul"]
keep = ("/attn/Wqkv/", "/attn/Wo/", "/mlp/Wi/") # int8; the MLP down-projection, head and rotary stay fp32
quantize_dynamic(str(fp32), str(int8), weight_type=QuantType.QInt8, per_channel=True,
nodes_to_exclude=[n for n in names if not any(k in n for k in keep)])
for p in (fp32, int8):
print(f"{p.name}: {p.stat().st_size / 1e6:.1f} MB")
if __name__ == "__main__":
ap = argparse.ArgumentParser()
ap.add_argument("dir")
ap.add_argument("--opset", type=int, default=17)
a = ap.parse_args()
export(Path(a.dir), a.opset)
|