From cd63de60a7e862239e40e640baa63ffb3ce6f4a8 Mon Sep 17 00:00:00 2001 From: Ivor Z Date: Mon, 12 Aug 2024 00:42:55 +0000 Subject: [PATCH] RKNN export Co-authored-by: Burhan Co-authored-by: Lakshantha Dissanayake --- ultralytics/cfg/__init__.py | 1 + ultralytics/engine/exporter.py | 21 ++++++++++++++------- ultralytics/nn/autobackend.py | 5 ++++- 3 files changed, 19 insertions(+), 8 deletions(-) diff --git a/ultralytics/cfg/__init__.py b/ultralytics/cfg/__init__.py index 4a23ec4267..9f250bc11f 100644 --- a/ultralytics/cfg/__init__.py +++ b/ultralytics/cfg/__init__.py @@ -183,6 +183,7 @@ CFG_BOOL_KEYS = { # boolean-only arguments "profile", "multi_scale", } +RKNN_CHIPS = {"rk3566", "rk3568", "rk3588", "rk3562", "rk3576", "rk2118"} def cfg2dict(cfg): diff --git a/ultralytics/engine/exporter.py b/ultralytics/engine/exporter.py index cc6fbfbc00..2cafc12cf7 100644 --- a/ultralytics/engine/exporter.py +++ b/ultralytics/engine/exporter.py @@ -64,7 +64,7 @@ from pathlib import Path import numpy as np import torch -from ultralytics.cfg import TASK2DATA, get_cfg +from ultralytics.cfg import TASK2DATA, get_cfg, RKNN_CHIPS from ultralytics.data import build_dataloader from ultralytics.data.dataset import YOLODataset from ultralytics.data.utils import check_cls_dataset, check_det_dataset @@ -112,6 +112,7 @@ def export_formats(): ["TensorFlow.js", "tfjs", "_web_model", True, False], ["PaddlePaddle", "paddle", "_paddle_model", True, True], ["NCNN", "ncnn", "_ncnn_model", True, True], + ["RKNN", "rknn", ".rknn", True, False], ] return pandas.DataFrame(x, columns=["Format", "Argument", "Suffix", "CPU", "GPU"]) @@ -183,7 +184,7 @@ 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, ncnn = flags # export booleans + jit, onnx, xml, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs, paddle, ncnn, rknn = flags # export booleans is_tf_format = any((saved_model, pb, tflite, edgetpu, tfjs)) # Device @@ -209,6 +210,8 @@ 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: + assert self.args.name in RKNN_CHIPS, f"Invalid processor name '{self.args.name}' for Rockchip RKNN export. Valid names are {RKNN_CHIPS}." if edgetpu: if not LINUX: raise SystemError("Edge TPU export only supported on Linux. See https://coral.ai/docs/edgetpu/compiler") @@ -323,6 +326,8 @@ class Exporter: f[10], _ = self.export_paddle() if ncnn: # NCNN f[11], _ = self.export_ncnn() + if rknn: #RKNN + f[12], _ = self.export_rknn() # Finish f = [str(x) for x in f if x] # filter out '' and None @@ -1022,14 +1027,16 @@ class Exporter: f, _ = self.export_onnx() from rknn.api import RKNN + platform = self.args.name # 'rk3566', 'rk3568', 'rk3588', 'rk3562', 'rk3576', 'rk2118' rknn = RKNN(verbose=False) - rknn.config(mean_values=[[0, 0, 0]], std_values=[ - [255, 255, 255]], target_platform='rk3588') - f = rknn.load_onnx(model=f) + rknn.config(mean_values=[[0, 0, 0]], std_values=[[255, 255, 255]], target_platform=platform) + # 'rk3566', 'rk3568', 'rk3588', 'rk3562', 'rk3576', 'rk2118' + # Must have quantization: 'rv1103', 'rv1106','rv1103b' + _ = rknn.load_onnx(model=f) #q = "int8" if self.args.int8 else "half" if self.args.half else "" # quantization #q = True if self.args.int8 else False - f = rknn.build(do_quantization=False) - f = rknn.export_rknn("yolov8n.rknn") + _ = rknn.build(do_quantization=False) + f = rknn.export_rknn(f.replace(".onnx", f"-{platform}.rknn")) LOGGER.info(f'\n{prefix} feed {f} to RKNN-Toolkit or RKNN-Toolkit2 to generate RKNN model.\n' 'Refer https://github.com/airockchip/rknn_model_zoo/tree/main/models/CV/object_detection/yolo') diff --git a/ultralytics/nn/autobackend.py b/ultralytics/nn/autobackend.py index cde35a5776..caa9c1df63 100644 --- a/ultralytics/nn/autobackend.py +++ b/ultralytics/nn/autobackend.py @@ -121,6 +121,7 @@ class AutoBackend(nn.Module): paddle, ncnn, triton, + rknn, ) = self._model_type(w) fp16 &= pt or jit or onnx or xml or engine or nn_module or triton # FP16 nhwc = coreml or saved_model or pb or tflite or edgetpu # BHWC formats (vs torch BCWH) @@ -551,7 +552,9 @@ class AutoBackend(nn.Module): ex.input(self.net.input_names()[0], mat_in) # WARNING: 'output_names' sorted as a temporary fix for https://github.com/pnnx/pnnx/issues/130 y = [np.array(ex.extract(x)[1])[None] for x in sorted(self.net.output_names())] - + elif getattr(self, 'rknn', False): + assert "for inference, please refer to https://github.com/airockchip/rknn_model_zoo/" + # NVIDIA Triton Inference Server elif self.triton: im = im.cpu().numpy() # torch to numpy