File size: 1,577 Bytes
1edc459 | 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 | import onnx
from onnxruntime.quantization import quantize_dynamic, QuantType
def export_onnx(
pt_path="models/yolo11-medium-page.pt",
onnx_path="models/yolo11-medium-page.onnx",
imgsz=736,
dynamic=False,
opset=12,
half=False,
):
from ultralytics import YOLO
print(f"Loading PT model: {pt_path}")
model = YOLO(pt_path)
print(f"Exporting to ONNX: {onnx_path}")
model.export(format="onnx", imgsz=imgsz, dynamic=dynamic, opset=opset, half=half)
return onnx_path
def downgrade_ir_v10_to_v9(input_path, output_path):
print(f"Downgrading IR v10 -> v9: {input_path} -> {output_path}")
model = onnx.load(input_path)
model.ir_version = 9
onnx.save(model, output_path)
return output_path
def quantize_model(input_path, output_path):
print(f"Quantizing model: {input_path} -> {output_path}")
quantize_dynamic(
model_input=input_path,
model_output=output_path,
weight_type=QuantType.QUInt8,
)
return output_path
def full_pipeline():
pt_path = "models/yolo11-medium-page-best.pt"
onnx_path = "models/yolo11-medium-page-best.onnx"
ir9_path = "models/yolo11-medium-page-best-ir9.onnx"
quant_path = "models/yolo11-medium-page-best-quant-ir9.onnx"
export_onnx(pt_path, onnx_path)
downgrade_ir_v10_to_v9(onnx_path, ir9_path)
quantize_model(ir9_path, quant_path)
import os
size_mb = os.path.getsize(quant_path) / 1024 / 1024
print(f"Done! Quantized model: {quant_path} ({size_mb:.1f} MB)")
if __name__ == "__main__":
full_pipeline()
|