From d6272f258e642bd318e09c943c7783d7bdc8d144 Mon Sep 17 00:00:00 2001 From: Laughing <61612323+Laughing-q@users.noreply.github.com> Date: Thu, 23 Jan 2025 20:28:41 +0800 Subject: [PATCH] RKNN export clean up (#18841) --- ultralytics/engine/exporter.py | 13 +++++-------- ultralytics/nn/autobackend.py | 2 +- 2 files changed, 6 insertions(+), 9 deletions(-) diff --git a/ultralytics/engine/exporter.py b/ultralytics/engine/exporter.py index cc18e179..ad2e6b07 100644 --- a/ultralytics/engine/exporter.py +++ b/ultralytics/engine/exporter.py @@ -1159,18 +1159,15 @@ class Exporter: 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}") + rknn.config(mean_values=[[0, 0, 0]], std_values=[[255, 255, 255]], target_platform=self.args.name) + rknn.load_onnx(model=f) + rknn.build(do_quantization=False) # TODO: Add quantization support + f = f.replace(".onnx", f"-{self.args.name}.rknn") + rknn.export_rknn(f"{export_path / f}") yaml_save(export_path / "metadata.yaml", self.metadata) return export_path, None diff --git a/ultralytics/nn/autobackend.py b/ultralytics/nn/autobackend.py index c979bc60..c63f999d 100644 --- a/ultralytics/nn/autobackend.py +++ b/ultralytics/nn/autobackend.py @@ -482,7 +482,7 @@ class AutoBackend(nn.Module): w = next(w.rglob("*.rknn")) # get *.rknn file from *_rknn_model dir rknn_model = RKNNLite() rknn_model.load_rknn(w) - ret = rknn_model.init_runtime() + rknn_model.init_runtime() metadata = Path(w).parent / "metadata.yaml" # Any other format (unsupported)