Aryan Sethi
Claude Opus 5 (1M context)
Quantise the ONNX export, and record how the model performs
35206ab | """Quantise an exported ONNX model to INT8, calibrating on real val frames. | |
| Static (not dynamic) quantisation, because this graph is almost entirely | |
| convolutions. Dynamic quantisation only converts operators whose weights can be | |
| quantised without knowing the activation range -- Linear and LSTM in practice -- | |
| so on a conv net it touches almost nothing and buys almost nothing. Static | |
| quantisation runs real images through the float model first, records the range | |
| each activation actually takes, and uses those ranges to pick scales. That | |
| calibration pass is why `CalibrationDataReader` exists. | |
| Two things that are easy to get wrong and silently produce a bad model: | |
| * **Preprocessing must match what the model was trained with.** The calibration | |
| data decides every activation scale, so feeding it differently-scaled images | |
| than inference will use makes every scale wrong. Here that means letterbox to | |
| 640 on a grey background, BGR to RGB, divide by 255, NCHW -- the same thing | |
| Ultralytics does at predict time. | |
| * **The graph needs `quant_pre_process` first.** Without shape inference and | |
| constant folding the quantiser leaves large parts of the graph in float, or | |
| fails outright, and the result is a model that is barely smaller and no | |
| faster. | |
| * **Activations must be unsigned on x64.** `QuantType.QInt8` for activations | |
| looks like the obvious symmetric choice and is 2.2x slower than | |
| `QUInt8` here -- 53.1 ms against 23.9 ms -- because onnxruntime's x64 INT8 | |
| kernels are written for unsigned activations and signed ones fall back to a | |
| slower path. onnxruntime warns about this and the warning is easy to miss. | |
| Weights stay signed; it is only the activations that matter. | |
| * **The detection head must stay in float.** Quantising the whole graph does | |
| not degrade this model, it destroys it -- measured here, full INT8 went from | |
| 213 detections to 0. The head decodes boxes with Sub/Add/Div into pixel | |
| coordinates spanning 0 to 640 and runs a Softmax over the DFL bins; one INT8 | |
| scale cannot carry that range and the decoded boxes collapse. The backbone, | |
| which is where the compute is, quantises fine. | |
| python onnx_quantization.py --model trained/yolo11n/weights/best.onnx | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| from pathlib import Path | |
| import cv2 | |
| import numpy as np | |
| from onnxruntime.quantization import CalibrationDataReader, QuantType, quantize_static | |
| from onnxruntime.quantization.shape_inference import quant_pre_process | |
| IMGSZ = 640 | |
| # How far back from the output counts as "the head". 12 covers the box decode, | |
| # the DFL softmax and the class sigmoid, and stops before the neck's convs -- | |
| # which are where the time actually goes, so excluding them would defeat the | |
| # exercise. | |
| HEAD_DEPTH = 12 | |
| def preprocess(path: str, size: int = IMGSZ) -> np.ndarray: | |
| """Image file -> the exact tensor the model sees at inference. | |
| Letterbox rather than a plain resize: squashing the aspect ratio would | |
| change the apparent height of every object, which is the quantity the depth | |
| estimator reads. | |
| """ | |
| image = cv2.imread(str(path)) | |
| if image is None: | |
| raise FileNotFoundError(path) | |
| height, width = image.shape[:2] | |
| scale = min(size / height, size / width) | |
| resized = cv2.resize(image, (int(round(width * scale)), int(round(height * scale)))) | |
| canvas = np.full((size, size, 3), 114, dtype=np.uint8) # Ultralytics' grey | |
| top = (size - resized.shape[0]) // 2 | |
| left = (size - resized.shape[1]) // 2 | |
| canvas[top:top + resized.shape[0], left:left + resized.shape[1]] = resized | |
| tensor = canvas[:, :, ::-1].astype(np.float32) / 255.0 # BGR -> RGB, 0..1 | |
| return np.ascontiguousarray(tensor.transpose(2, 0, 1))[None] # NCHW, batch 1 | |
| class ValReader(CalibrationDataReader): | |
| """Feeds calibration images one at a time until exhausted.""" | |
| def __init__(self, paths, input_name: str = "images", n: int = 300): | |
| self.data = iter([{input_name: preprocess(p)} for p in paths[:n]]) | |
| def get_next(self): | |
| return next(self.data, None) | |
| def val_image_paths(root: Path, dataset: str, split: str, n: int) -> list[str]: | |
| """Calibration frames, taken from the split the model is scored on. | |
| Sorted rather than sampled so a rebuild calibrates on the same images, and | |
| drawn from val rather than train because the activation ranges that matter | |
| are the ones seen at inference. | |
| """ | |
| from src.common import paths as project_paths | |
| from src.common import schema | |
| manifest = schema.read_manifest(root / "manifest.parquet") | |
| manifest = manifest[(manifest["source"] == dataset) & (manifest["split"] == split)] | |
| images = sorted(manifest["image_path"].unique())[:n] | |
| return [str(project_paths.resolve_image(root, p)) for p in images] | |
| def head_nodes(model_path: str, depth: int = HEAD_DEPTH) -> list[str]: | |
| """Names of the nodes within `depth` steps of the graph output. | |
| Walked backwards from the output rather than matched by name, because node | |
| names are an export artefact and change between opsets and Ultralytics | |
| versions, while graph distance from the output does not. | |
| """ | |
| import onnx | |
| graph = onnx.load(model_path, load_external_data=False).graph | |
| producer = {out: node for node in graph.node for out in node.output} | |
| names, seen = [], {graph.output[0].name: 0} | |
| frontier = [graph.output[0].name] | |
| while frontier: | |
| tensor = frontier.pop(0) | |
| node = producer.get(tensor) | |
| if node is None or seen[tensor] > depth: | |
| continue | |
| names.append(node.name) | |
| for inp in node.input: | |
| if inp not in seen: | |
| seen[inp] = seen[tensor] + 1 | |
| frontier.append(inp) | |
| return names | |
| def main() -> None: | |
| parser = argparse.ArgumentParser(description=__doc__, | |
| formatter_class=argparse.RawDescriptionHelpFormatter) | |
| parser.add_argument("--model", default="trained/yolo11n/weights/best.onnx") | |
| parser.add_argument("--out", default=None, help="defaults to <model>_int8.onnx") | |
| parser.add_argument("--unified-root", type=Path, default=Path("data/unified")) | |
| parser.add_argument("--dataset", default="nuscenes") | |
| parser.add_argument("--split", default="val") | |
| parser.add_argument("--calib-images", type=int, default=300) | |
| parser.add_argument("--head-depth", type=int, default=HEAD_DEPTH, | |
| help="nodes this close to the output stay float32. " | |
| "0 quantises everything, which produces a model " | |
| "that loads and detects nothing.") | |
| parser.add_argument("--per-tensor", action="store_true", | |
| help="one weight scale per tensor instead of per output " | |
| "channel. Smaller, and markedly less accurate.") | |
| args = parser.parse_args() | |
| model = Path(args.model) | |
| output = Path(args.out) if args.out else model.with_name(model.stem + "_int8.onnx") | |
| prepared = model.with_name(model.stem + "_prepared.onnx") | |
| print(f"preparing {model} (shape inference + constant folding)") | |
| quant_pre_process(str(model), str(prepared), skip_symbolic_shape=True) | |
| exclude = head_nodes(str(prepared), args.head_depth) if args.head_depth else [] | |
| print(f"keeping {len(exclude)} head nodes in float32") | |
| images = val_image_paths(args.unified_root, args.dataset, args.split, | |
| args.calib_images) | |
| print(f"calibrating on {len(images)} {args.dataset} {args.split} frames") | |
| quantize_static( | |
| str(prepared), str(output), | |
| ValReader(images, n=args.calib_images), | |
| # Unsigned activations, signed weights. See the note at the top: the | |
| # x64 kernels are built for this combination and the obvious-looking | |
| # QInt8/QInt8 is more than twice as slow. | |
| activation_type=QuantType.QUInt8, | |
| weight_type=QuantType.QInt8, | |
| # Per output channel: a convolution's filters have very different | |
| # magnitudes, and one scale across all of them wastes most of the INT8 | |
| # range on the largest filter. | |
| per_channel=not args.per_tensor, | |
| nodes_to_exclude=exclude, | |
| ) | |
| prepared.unlink(missing_ok=True) | |
| before = model.stat().st_size + sum( | |
| f.stat().st_size for f in model.parent.glob(model.name + ".data")) | |
| after = output.stat().st_size | |
| print(f"\n{model.name} {before / 1e6:.1f} MB") | |
| print(f"{output.name} {after / 1e6:.1f} MB ({before / after:.2f}x smaller)") | |
| if __name__ == "__main__": | |
| main() | |