ultralytics 8.3.65 Rockchip RKNN Integration for Ultralytics YOLO models (#16308)
Signed-off-by: Francesco Mattioli <Francesco.mttl@gmail.com> Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com> Co-authored-by: Burhan <62214284+Burhan-Q@users.noreply.github.com> Co-authored-by: Lakshantha Dissanayake <lakshantha@ultralytics.com> Co-authored-by: Burhan <Burhan-Q@users.noreply.github.com> Co-authored-by: Laughing-q <1185102784@qq.com> Co-authored-by: UltralyticsAssistant <web@ultralytics.com> Co-authored-by: Laughing <61612323+Laughing-q@users.noreply.github.com> Co-authored-by: Ultralytics Assistant <135830346+UltralyticsAssistant@users.noreply.github.com> Co-authored-by: Lakshantha Dissanayake <lakshanthad@yahoo.com> Co-authored-by: Francesco Mattioli <Francesco.mttl@gmail.com> Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
Ivor Zhu
GitHub
parent
617dea8e25
commit
b5e0cee943
41 changed files with 390 additions and 118 deletions
|
|
@ -19,6 +19,7 @@ PaddlePaddle | `paddle` | yolo11n_paddle_model/
|
|||
MNN | `mnn` | yolo11n.mnn
|
||||
NCNN | `ncnn` | yolo11n_ncnn_model/
|
||||
IMX | `imx` | yolo11n_imx_model/
|
||||
RKNN | `rknn` | yolo11n_rknn_model/
|
||||
|
||||
Requirements:
|
||||
$ pip install "ultralytics[export]"
|
||||
|
|
@ -78,11 +79,13 @@ from ultralytics.nn.tasks import DetectionModel, SegmentationModel, WorldModel
|
|||
from ultralytics.utils import (
|
||||
ARM64,
|
||||
DEFAULT_CFG,
|
||||
IS_COLAB,
|
||||
IS_JETSON,
|
||||
LINUX,
|
||||
LOGGER,
|
||||
MACOS,
|
||||
PYTHON_VERSION,
|
||||
RKNN_CHIPS,
|
||||
ROOT,
|
||||
WINDOWS,
|
||||
__version__,
|
||||
|
|
@ -122,6 +125,7 @@ def export_formats():
|
|||
["MNN", "mnn", ".mnn", True, True, ["batch", "half", "int8"]],
|
||||
["NCNN", "ncnn", "_ncnn_model", True, True, ["batch", "half"]],
|
||||
["IMX", "imx", "_imx_model", True, True, ["int8"]],
|
||||
["RKNN", "rknn", "_rknn_model", False, False, ["batch", "name"]],
|
||||
]
|
||||
return dict(zip(["Format", "Argument", "Suffix", "CPU", "GPU", "Arguments"], zip(*x)))
|
||||
|
||||
|
|
@ -226,22 +230,10 @@ class Exporter:
|
|||
flags = [x == fmt for x in fmts]
|
||||
if sum(flags) != 1:
|
||||
raise ValueError(f"Invalid export format='{fmt}'. Valid formats are {fmts}")
|
||||
(
|
||||
jit,
|
||||
onnx,
|
||||
xml,
|
||||
engine,
|
||||
coreml,
|
||||
saved_model,
|
||||
pb,
|
||||
tflite,
|
||||
edgetpu,
|
||||
tfjs,
|
||||
paddle,
|
||||
mnn,
|
||||
ncnn,
|
||||
imx,
|
||||
) = flags # export booleans
|
||||
(jit, onnx, xml, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs, paddle, mnn, ncnn, imx, rknn) = (
|
||||
flags # export booleans
|
||||
)
|
||||
|
||||
is_tf_format = any((saved_model, pb, tflite, edgetpu, tfjs))
|
||||
|
||||
# Device
|
||||
|
|
@ -277,6 +269,16 @@ class Exporter:
|
|||
if self.args.optimize:
|
||||
assert not ncnn, "optimize=True not compatible with format='ncnn', i.e. use optimize=False"
|
||||
assert self.device.type == "cpu", "optimize=True not compatible with cuda devices, i.e. use device='cpu'"
|
||||
if rknn:
|
||||
if not self.args.name:
|
||||
LOGGER.warning(
|
||||
"WARNING ⚠️ Rockchip RKNN export requires a missing 'name' arg for processor type. Using default name='rk3588'."
|
||||
)
|
||||
self.args.name = "rk3588"
|
||||
self.args.name = self.args.name.lower()
|
||||
assert self.args.name in RKNN_CHIPS, (
|
||||
f"Invalid processor name '{self.args.name}' for Rockchip RKNN export. Valid names are {RKNN_CHIPS}."
|
||||
)
|
||||
if self.args.int8 and tflite:
|
||||
assert not getattr(model, "end2end", False), "TFLite INT8 export not supported for end2end models."
|
||||
if edgetpu:
|
||||
|
|
@ -417,6 +419,8 @@ class Exporter:
|
|||
f[12], _ = self.export_ncnn()
|
||||
if imx:
|
||||
f[13], _ = self.export_imx()
|
||||
if rknn:
|
||||
f[14], _ = self.export_rknn()
|
||||
|
||||
# Finish
|
||||
f = [str(x) for x in f if x] # filter out '' and None
|
||||
|
|
@ -746,7 +750,7 @@ class Exporter:
|
|||
model = IOSDetectModel(self.model, self.im) if self.args.nms else self.model
|
||||
else:
|
||||
if self.args.nms:
|
||||
LOGGER.warning(f"{prefix} WARNING ⚠️ 'nms=True' is only available for Detect models like 'yolov8n.pt'.")
|
||||
LOGGER.warning(f"{prefix} WARNING ⚠️ 'nms=True' is only available for Detect models like 'yolo11n.pt'.")
|
||||
# TODO CoreML Segment and Pose model pipelining
|
||||
model = self.model
|
||||
|
||||
|
|
@ -1141,6 +1145,35 @@ class Exporter:
|
|||
return f, None
|
||||
|
||||
@try_export
|
||||
def export_rknn(self, prefix=colorstr("RKNN:")):
|
||||
"""YOLO RKNN model export."""
|
||||
LOGGER.info(f"\n{prefix} starting export with rknn-toolkit2...")
|
||||
|
||||
check_requirements("rknn-toolkit2")
|
||||
if IS_COLAB:
|
||||
# Prevent 'exit' from closing the notebook https://github.com/airockchip/rknn-toolkit2/issues/259
|
||||
import builtins
|
||||
|
||||
builtins.exit = lambda: None
|
||||
|
||||
from rknn.api import RKNN
|
||||
|
||||
f, _ = self.export_onnx()
|
||||
|
||||
platform = self.args.name
|
||||
|
||||
export_path = Path(f"{Path(f).stem}_rknn_model")
|
||||
export_path.mkdir(exist_ok=True)
|
||||
|
||||
rknn = RKNN(verbose=False)
|
||||
rknn.config(mean_values=[[0, 0, 0]], std_values=[[255, 255, 255]], target_platform=platform)
|
||||
_ = rknn.load_onnx(model=f)
|
||||
_ = rknn.build(do_quantization=False) # TODO: Add quantization support
|
||||
f = f.replace(".onnx", f"-{platform}.rknn")
|
||||
_ = rknn.export_rknn(f"{export_path / f}")
|
||||
yaml_save(export_path / "metadata.yaml", self.metadata)
|
||||
return export_path, None
|
||||
|
||||
def export_imx(self, prefix=colorstr("IMX:")):
|
||||
"""YOLO IMX export."""
|
||||
gptq = False
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue