Fabio
c36510c21e
Some checks failed
Linux Build / Linux cpu (push) Has been cancelled
Linux Build / Linux cpu-clang (push) Has been cancelled
Linux Build / Linux vulkan (push) Has been cancelled
mac-build / macOS CPU (push) Has been cancelled
Nix Build / Nix (cpu) (push) Has been cancelled
Nix Build / Nix (vulkan) (push) Has been cancelled
windows-build / Windows CPU (push) Has been cancelled
73 lines
2.9 KiB
Python
73 lines
2.9 KiB
Python
#!/usr/bin/env python3
|
|
"""Engine-build plumbing shared by the `tools/rknn/prepare_*` converters.
|
|
|
|
Every converter ends in the same five calls on the RKNPU2 Python API, and the compiler-log gate
|
|
(`rknn_build_log.rknn_build_gate`) has to wrap `rknn.build` in every one of them, so the sequence
|
|
lives here. Everything that produces its arguments stays in the converter: the ONNX surgery, the
|
|
input names and shapes, the per-engine file naming, and the simulator-side verification.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import sys
|
|
from pathlib import Path
|
|
from typing import Optional, Sequence
|
|
|
|
|
|
def report_engine(output, *details: str) -> None:
|
|
"""One summary line per written artefact, the same everywhere: ``wrote <path> (<size> MB)``
|
|
plus any per-converter details the caller wants on the line. Scripts parse nothing off
|
|
these; they exist so the five converters report one way."""
|
|
output = Path(output)
|
|
detail = "".join(f", {d}" for d in details)
|
|
print(f"wrote {output} ({output.stat().st_size / 1e6:.1f} MB{detail})")
|
|
|
|
|
|
def build_engine(
|
|
onnx_path: Path,
|
|
output: Path,
|
|
target: str,
|
|
*,
|
|
names: Optional[Sequence[str]] = None,
|
|
shapes: Optional[Sequence[Sequence[int]]] = None,
|
|
opt_level: Optional[int] = None,
|
|
verbose: bool = False,
|
|
gate_label: Optional[str] = None,
|
|
gate_log: Optional[Path] = None,
|
|
) -> None:
|
|
"""Compile ``onnx_path`` into ``output`` for ``target``: config, load, gated build, export.
|
|
|
|
``names``/``shapes`` are declared to ``load_onnx`` only for a graph that is still dynamic
|
|
there. A converter that cropped the sequence axis into a fixed window compiles a static graph
|
|
and declares nothing: re-declaring shapes on a static graph only warns.
|
|
|
|
Not ``single_core_mode``: the engine stays multi-core, and ``rknn_set_core_mask`` in the
|
|
runtime decides how many NPU cores a call actually runs on.
|
|
"""
|
|
from rknn.api import RKNN
|
|
|
|
output = Path(output)
|
|
onnx_path = Path(onnx_path)
|
|
|
|
rknn = RKNN(verbose=verbose)
|
|
config: dict = {"target_platform": target}
|
|
if opt_level is not None:
|
|
config["optimization_level"] = opt_level
|
|
rknn.config(**config)
|
|
if names is not None and shapes is not None:
|
|
ret = rknn.load_onnx(model=str(onnx_path), inputs=list(names),
|
|
input_size_list=[list(s) for s in shapes])
|
|
else:
|
|
ret = rknn.load_onnx(model=str(onnx_path))
|
|
if ret != 0:
|
|
raise SystemExit("load_onnx failed")
|
|
sys.path.insert(0, str(Path(__file__).resolve().parent))
|
|
from rknn_build_log import rknn_build_gate
|
|
|
|
with rknn_build_gate(gate_label or output.name, gate_log or output.with_suffix(".build.log")):
|
|
built = rknn.build(do_quantization=False)
|
|
if built != 0:
|
|
raise SystemExit("build failed (see the rknn log above)")
|
|
if rknn.export_rknn(str(output)) != 0:
|
|
raise SystemExit("export_rknn failed")
|
|
rknn.release()
|