File size: 4,708 Bytes
35206ab | 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 | """Quantise the PyTorch checkpoint, and measure whether it was worth it.
The short answer is that it is not, and the reason is worth knowing rather than
taking on trust -- so this quantises the model and benchmarks the result instead
of asserting anything.
PyTorch offers two routes on CPU:
* **Dynamic quantisation** is one line and needs no calibration, but it only
converts operators whose weights can be quantised without knowing the
activation range: Linear, LSTM, GRU, Embedding. YOLO11n is convolutions --
81 Conv and 7 ConvolutionDepthWise in the exported graph, and a handful of
Linear layers at most. So it converts almost nothing.
* **Static quantisation** does cover convolutions, but requires the model to be
built for it: QuantStub and DeQuantStub around the graph, fusable Conv-BN-ReLU
patterns, and a forward that traces cleanly. Ultralytics' model is none of
those things, and making it so means forking the architecture.
Which is why the deployment path for this project is export-then-quantise --
ONNX, NCNN, TFLite -- rather than quantising the checkpoint. The exporters
lower the graph to plain convolutions first, and the runtimes have kernels that
actually execute INT8 convolutions quickly. `onnx_quantization.py` is that path,
and it gets 4.5x smaller at no measurable accuracy cost.
python pt_quantization.py --weights trained/yolo11n/weights/best.pt
"""
from __future__ import annotations
import argparse
import time
from pathlib import Path
import numpy as np
import torch
def count_layers(module: torch.nn.Module) -> dict[str, int]:
counts: dict[str, int] = {}
for layer in module.modules():
name = type(layer).__name__
if name in ("Conv2d", "Linear", "BatchNorm2d", "ConvTranspose2d"):
counts[name] = counts.get(name, 0) + 1
return counts
def time_forward(model: torch.nn.Module, tensor: torch.Tensor,
runs: int, warmup: int) -> tuple[float, float]:
"""p50 and p95 milliseconds for a raw forward pass.
Raw forward rather than Ultralytics' predict(), so the measurement is the
model and not the pre/post-processing around it -- those are identical
between the two variants and would dilute the comparison.
"""
with torch.inference_mode():
for _ in range(warmup):
model(tensor)
samples = []
for _ in range(runs):
started = time.perf_counter()
model(tensor)
samples.append((time.perf_counter() - started) * 1000.0)
values = np.array(samples)
return float(np.percentile(values, 50)), float(np.percentile(values, 95))
def main() -> None:
parser = argparse.ArgumentParser(description=__doc__,
formatter_class=argparse.RawDescriptionHelpFormatter)
parser.add_argument("--weights", default="trained/yolo11n/weights/best.pt")
parser.add_argument("--imgsz", type=int, default=640)
parser.add_argument("--runs", type=int, default=30)
parser.add_argument("--warmup", type=int, default=5)
parser.add_argument("--threads", type=int, default=4)
args = parser.parse_args()
from ultralytics import YOLO
torch.set_num_threads(args.threads)
model = YOLO(args.weights).model.float().eval()
tensor = torch.zeros(1, 3, args.imgsz, args.imgsz)
print(f"layers: {count_layers(model)}")
quantised = torch.ao.quantization.quantize_dynamic(
model, {torch.nn.Linear, torch.nn.LSTM, torch.nn.GRU}, dtype=torch.qint8)
converted = sum(1 for layer in quantised.modules()
if "quantized" in type(layer).__module__)
print(f"layers actually converted to int8: {converted}")
original = Path(args.weights).stat().st_size
out = Path(args.weights).with_name(Path(args.weights).stem + "_dynint8.pt")
torch.save(quantised.state_dict(), out)
print(f"\nsize fp32 {original / 1e6:5.1f} MB "
f"dynamic int8 {out.stat().st_size / 1e6:5.1f} MB")
fp32 = time_forward(model, tensor, args.runs, args.warmup)
int8 = time_forward(quantised, tensor, args.runs, args.warmup)
print(f"\nforward pass at {args.imgsz}px, {args.threads} threads, "
f"{args.runs} runs")
print(f" fp32 p50 {fp32[0]:7.2f} ms p95 {fp32[1]:7.2f} ms")
print(f" dynamic int8 p50 {int8[0]:7.2f} ms p95 {int8[1]:7.2f} ms")
print(f" speedup {fp32[0] / int8[0]:.2f}x")
if converted == 0:
print("\nNothing was converted. Dynamic quantisation covers Linear, LSTM,")
print("GRU and Embedding; this graph is convolutions. Use the export")
print("path instead -- see onnx_quantization.py.")
if __name__ == "__main__":
main()
|