File size: 3,417 Bytes
ebeed67 | 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 | """The heavy export steps, each run in its own child process.
``python -m camstack_train._steps <step> '<json args>'``. A crash inside
coremltools, NNCF or a converter takes down that child only; the parent
records the step as failed (``export.SubprocessToolchain``). The recipes mirror
``scripts/build-camstack-models.py`` / ``convert_lib.py`` in camstack/server
(the yolov9 recipe: static ONNX at opset 13, ``nms=False``, batch 1; CoreML
``half=True`` — FP16, D743; OpenVINO from the ONNX via ``ov.convert_model``,
never the Ultralytics OV exporter), and are not imported from them.
"""
from __future__ import annotations
import json
import shutil
import sys
from pathlib import Path
#: Same bound as convert_lib / the inference pool (CAMSTACK_OV_COMPILATION_THREADS).
OV_COMPILATION_THREADS = 4
ONNX_OPSET = 13
def _move(src: Path, dest: Path) -> None:
if dest.exists():
shutil.rmtree(dest) if dest.is_dir() else dest.unlink()
dest.parent.mkdir(parents=True, exist_ok=True)
shutil.move(str(src), str(dest))
def step_onnx(pt: str, out: str, imgsz: int) -> None:
from ultralytics import YOLO
path = YOLO(pt).export(
format="onnx", imgsz=imgsz, dynamic=False, simplify=True,
opset=ONNX_OPSET, nms=False, batch=1, half=False,
)
_move(Path(path), Path(out))
def step_coreml(pt: str, out: str, imgsz: int) -> None:
from ultralytics import YOLO
path = YOLO(pt).export(format="coreml", imgsz=imgsz, dynamic=False, batch=1, half=True, nms=False)
_move(Path(path), Path(out))
def _validate_ir(xml: str) -> None:
import openvino as ov
core = ov.Core()
compiled = core.compile_model(core.read_model(xml), "CPU", {"COMPILATION_NUM_THREADS": OV_COMPILATION_THREADS})
for port in list(compiled.inputs) + list(compiled.outputs):
if not port.get_any_name():
raise ValueError("the IR has an unnamed tensor")
def step_openvino(onnx: str, xml: str, fp16: bool) -> None:
import openvino as ov
Path(xml).parent.mkdir(parents=True, exist_ok=True)
ov.save_model(ov.convert_model(onnx), xml, compress_to_fp16=fp16)
_validate_ir(xml)
def step_openvino_int8(onnx: str, xml: str, calib: list[str], imgsz: int) -> None:
import nncf
import numpy as np
import openvino as ov
from ._vendor import ensure_pipeline_on_path
ensure_pipeline_on_path()
import inference_pool # the pool's own decode + letterbox (grey 114, /255)
def preprocess(path: str):
arr, _scale, _pad = inference_pool.letterbox(inference_pool.decode_jpeg(Path(path).read_bytes()), imgsz)
return arr.transpose(2, 0, 1)[np.newaxis].astype(np.float32)
dataset = nncf.Dataset(calib, preprocess)
quantized = nncf.quantize(ov.convert_model(onnx), dataset, subset_size=len(calib))
Path(xml).parent.mkdir(parents=True, exist_ok=True)
ov.save_model(quantized, xml)
_validate_ir(xml)
STEPS = {
"onnx": lambda a: step_onnx(a["pt"], a["out"], a["imgsz"]),
"coreml": lambda a: step_coreml(a["pt"], a["out"], a["imgsz"]),
"openvino": lambda a: step_openvino(a["onnx"], a["xml"], a["fp16"]),
"openvino-int8": lambda a: step_openvino_int8(a["onnx"], a["xml"], a["calib"], a["imgsz"]),
}
def main(argv: list[str]) -> int:
step, payload = argv[0], json.loads(argv[1])
STEPS[step](payload)
return 0
if __name__ == "__main__":
sys.exit(main(sys.argv[1:]))
|