HongChuang/ultralytics-ascend
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 3

RKNN export clean up (#18841)

Browse source
This commit is contained in:
Laughing 2025-01-23 20:28:41 +08:00 • committed by GitHub
parent 729ef7f270
commit d6272f258e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 6 additions and 9 deletions
Download patch file Download diff file Expand all files Collapse all files

13
ultralytics/engine/exporter.py
View file

@ -1159,18 +1159,15 @@ class Exporter:
from rknn.api import RKNN from rknn.api import RKNN
f, _ = self.export_onnx() f, _ = self.export_onnx()
platform = self.args.name
export_path = Path(f"{Path(f).stem}_rknn_model") export_path = Path(f"{Path(f).stem}_rknn_model")
export_path.mkdir(exist_ok=True) export_path.mkdir(exist_ok=True)
rknn = RKNN(verbose=False) rknn = RKNN(verbose=False)
rknn.config(mean_values=[[0, 0, 0]], std_values=[[255, 255, 255]], target_platform=platform) rknn.config(mean_values=[[0, 0, 0]], std_values=[[255, 255, 255]], target_platform=self.args.name)
_ = rknn.load_onnx(model=f) rknn.load_onnx(model=f)
_ = rknn.build(do_quantization=False) # TODO: Add quantization support rknn.build(do_quantization=False) # TODO: Add quantization support
f = f.replace(".onnx", f"-{platform}.rknn") f = f.replace(".onnx", f"-{self.args.name}.rknn")
_ = rknn.export_rknn(f"{export_path / f}") rknn.export_rknn(f"{export_path / f}")
yaml_save(export_path / "metadata.yaml", self.metadata) yaml_save(export_path / "metadata.yaml", self.metadata)
return export_path, None return export_path, None

2
ultralytics/nn/autobackend.py
View file

@ -482,7 +482,7 @@ class AutoBackend(nn.Module):
w = next(w.rglob("*.rknn")) # get *.rknn file from *_rknn_model dir w = next(w.rglob("*.rknn")) # get *.rknn file from *_rknn_model dir
rknn_model = RKNNLite() rknn_model = RKNNLite()
rknn_model.load_rknn(w) rknn_model.load_rknn(w)
ret = rknn_model.init_runtime() rknn_model.init_runtime()
metadata = Path(w).parent / "metadata.yaml" metadata = Path(w).parent / "metadata.yaml"
# Any other format (unsupported) # Any other format (unsupported)