From c72cf61ed00142a4e3b451a05fc5acfd8cef0a7e Mon Sep 17 00:00:00 2001 From: yoni13 Date: Sun, 12 Jan 2025 01:24:24 +0800 Subject: [PATCH] support core_mask for specfic socs --- machine-learning/rknn/rknnpool.py | 21 +++++++++++++++++++-- 1 file changed, 19 insertions(+), 2 deletions(-) diff --git a/machine-learning/rknn/rknnpool.py b/machine-learning/rknn/rknnpool.py index 1bb967904d..6787fbbf87 100644 --- a/machine-learning/rknn/rknnpool.py +++ b/machine-learning/rknn/rknnpool.py @@ -5,7 +5,8 @@ import os from concurrent.futures import ThreadPoolExecutor from queue import Queue -supported_socs = ["rk3566", "rk3568", "rk3588"] +supported_socs = ["rk3562", "rk3566", "rk3568", "rk3576", "rk3588"] +coremask_supported_socs = ["rk3576","rk3588"] try: from rknnlite.api import RKNNLite @@ -15,6 +16,7 @@ try: for soc in supported_socs: if soc in device_compatible_str: is_available = True + soc_name = soc break else: is_available = False @@ -29,7 +31,22 @@ def initRKNN(rknnModel="./rknnModel/yolov5s.rknn", id=0): if ret != 0: print("Load RKNN rknnModel failed") exit(ret) - ret = rknn_lite.init_runtime() + + if soc_name in coremask_supported_socs: + if id == 0: + core_mask = RKNNLite.NPU_CORE_0 + elif id == 1: + core_mask = RKNNLite.NPU_CORE_1 + elif id == 2: + core_mask = RKNNLite.NPU_CORE_2 + elif id == -1: + core_mask = RKNNLite.NPU_CORE_0_1_2 + else: + core_mask = None + else: + core_mask = None + + ret = rknn_lite.init_runtime(core_mask=core_mask) if ret != 0: print("Init runtime environment failed") exit(ret)