"""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 _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()