File size: 5,608 Bytes
33bfb6d | 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 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 | # /// script
# requires-python = ">=3.11,<3.14"
# dependencies = ["onnx>=1.17", "numpy<3", "onnxruntime>=1.20"]
# ///
"""Shrinks a float ONNX model for download without changing how it runs.
The browser computes in float32 (onnxruntime-web's int8 kernels are several times slower than its float ones), so
only the download is compressed: every large weight is stored as int8 with a float32 scale per block of 64 values.
The page expands them back to float32 and hands them to onnxruntime-web as the model's external data.
Writes, next to each other:
graph.onnx the model graph; large weights point into "weights.f32" (external data)
weights.q8 int8 values, then float32 scales, tensor by tensor
weights.json for each tensor: where its int8 values and scales start, its shape, and where it goes in weights.f32
uv run compress.py v2-release/browser/model.onnx v2-release/web
"""
import json
import sys
from pathlib import Path
import numpy as np
import onnx
from onnx import numpy_helper
from onnx.external_data_helper import set_external_data
BLOCK = 64
MIN_ELEMENTS = 4096 # smaller tensors (norm weights) stay inside the graph as float32
EXTERNAL = "weights.f32"
def quantize(values: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
"""Symmetric int8 with one scale per BLOCK consecutive values."""
flat = values.astype(np.float32).ravel()
padded = np.pad(flat, (0, (-len(flat)) % BLOCK)).reshape(-1, BLOCK)
scales = np.abs(padded).max(axis=1) / 127
scales[scales == 0] = 1
q = np.clip(np.round(padded / scales[:, None]), -127, 127).astype(np.int8)
return q.ravel()[: len(flat)], scales.astype(np.float32)
def dequantize(q: np.ndarray, scales: np.ndarray, count: int) -> np.ndarray:
padded = np.pad(q.astype(np.float32), (0, (-count) % BLOCK)).reshape(-1, BLOCK)
return (padded * scales[:, None]).ravel()[:count]
def dedupe_transposes(model) -> None:
"""The output layer reuses the embedding table, but export stores a transposed copy. Keep one and transpose it
in the graph instead, which saves about 2.6 MB of download."""
from onnx import helper
graph = model.graph
arrays = {t.name: numpy_helper.to_array(t) for t in graph.initializer if t.dims and np.prod(t.dims) >= MIN_ELEMENTS}
for name, values in list(arrays.items()):
if values.ndim != 2:
continue
for other, candidate in arrays.items():
if other != name and candidate.shape == values.T.shape and np.array_equal(candidate, values.T):
transposed = f"{other}_transposed"
graph.node.insert(0, helper.make_node("Transpose", [other], [transposed], perm=[1, 0], name=transposed))
for node in graph.node:
node.input[:] = [transposed if i == name else i for i in node.input]
graph.initializer.remove(next(t for t in graph.initializer if t.name == name))
print(f"stored {name} as a transpose of {other}")
return
def compress(src: Path, out: Path) -> dict:
out.mkdir(parents=True, exist_ok=True)
model = onnx.load(str(src))
dedupe_transposes(model)
manifest, q8 = [], bytearray()
f32_offset = 0
for tensor in model.graph.initializer:
values = numpy_helper.to_array(tensor)
if values.dtype != np.float32 or values.size < MIN_ELEMENTS:
continue
q, scales = quantize(values)
entry = {"name": tensor.name, "shape": list(values.shape), "count": int(values.size), "f32_offset": f32_offset,
"q_offset": len(q8), "scale_offset": len(q8) + len(q) + (-len(q)) % 4}
q8 += q.tobytes()
q8 += b"\0" * ((-len(q8)) % 4) # keep scales 4-byte aligned for Float32Array views
q8 += scales.tobytes()
manifest.append(entry)
# Point the graph at weights.f32, which the browser builds by expanding weights.q8.
if not tensor.HasField("raw_data"):
tensor.raw_data = values.astype(np.float32).tobytes()
tensor.ClearField("float_data")
set_external_data(tensor, location=EXTERNAL, offset=f32_offset, length=values.size * 4)
tensor.ClearField("raw_data")
tensor.data_location = onnx.TensorProto.EXTERNAL
f32_offset += values.size * 4
onnx.save_model(model, str(out / "graph.onnx"))
(out / "weights.q8").write_bytes(bytes(q8))
(out / "weights.json").write_text(json.dumps({"block": BLOCK, "f32_bytes": f32_offset, "tensors": manifest}))
return {"graph_mb": round((out / "graph.onnx").stat().st_size / 1e6, 3), "weights_q8_mb": round(len(q8) / 1e6, 2),
"expanded_mb": round(f32_offset / 1e6, 2), "tensors": len(manifest)}
def expand(out: Path) -> Path:
"""What the browser does: rebuild weights.f32 from weights.q8. Used here to check accuracy."""
meta = json.loads((out / "weights.json").read_text())
raw = (out / "weights.q8").read_bytes()
f32 = np.zeros(meta["f32_bytes"] // 4, np.float32)
for t in meta["tensors"]:
q = np.frombuffer(raw, np.int8, t["count"], t["q_offset"])
blocks = (t["count"] + meta["block"] - 1) // meta["block"]
scales = np.frombuffer(raw, np.float32, blocks, t["scale_offset"])
f32[t["f32_offset"] // 4: t["f32_offset"] // 4 + t["count"]] = dequantize(q, scales, t["count"])
path = out / EXTERNAL
f32.tofile(path)
return path
if __name__ == "__main__":
src, out = Path(sys.argv[1]), Path(sys.argv[2])
print(json.dumps(compress(src, out), indent=1))
expand(out)
print(f"expanded {out / EXTERNAL} for checking")
|