cfahlgren1 HF Staff
v2: 13 MB browser build (int8 download, float compute), export/compress scripts, onnx evals
33bfb6d verified | # /// 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") | |