From 6e08f1f37147a651caad6e0b60fbe18843d07c23 Mon Sep 17 00:00:00 2001 From: yoni13 Date: Sat, 25 Jan 2025 19:14:34 +0800 Subject: [PATCH] dont stuck on core_0 on rk3588 but untested --- machine-learning/rknn/rknnpool.py | 15 +++------------ 1 file changed, 3 insertions(+), 12 deletions(-) diff --git a/machine-learning/rknn/rknnpool.py b/machine-learning/rknn/rknnpool.py index 17c0db021f..14fa06481e 100644 --- a/machine-learning/rknn/rknnpool.py +++ b/machine-learning/rknn/rknnpool.py @@ -32,7 +32,7 @@ except (FileNotFoundError, ImportError): soc_name = None -def init_rknn(rknnModel, id) -> Callable: +def init_rknn(rknnModel) -> Callable: if not is_available: raise RuntimeError("rknn is not available!") rknn_lite = RKNNLite() @@ -41,16 +41,7 @@ def init_rknn(rknnModel, id) -> Callable: raise RuntimeError("Load RKNN rknnModel failed") if soc_name in coremask_supported_socs: - if id == 0: - ret = rknn_lite.init_runtime(core_mask=RKNNLite.NPU_CORE_0) - elif id == 1: - ret = rknn_lite.init_runtime(core_mask=RKNNLite.NPU_CORE_1) - elif id == 2: - ret = rknn_lite.init_runtime(core_mask=RKNNLite.NPU_CORE_2) - elif id == -1: - ret = rknn_lite.init_runtime(core_mask=RKNNLite.NPU_CORE_0_1_2) - else: - ret = rknn_lite.init_runtime() + ret = rknn_lite.init_runtime(core_mask=RKNNLite.NPU_CORE_AUTO) else: ret = rknn_lite.init_runtime() # Please do not set this parameter on other platforms. @@ -63,7 +54,7 @@ def init_rknn(rknnModel, id) -> Callable: def init_rknns(rknnModel, tpes) -> list[Callable]: rknn_list = [] for i in range(tpes): - rknn_list.append(init_rknn(rknnModel, i % 3)) + rknn_list.append(init_rknn(rknnModel)) return rknn_list