[Fix] fix rknn quantization pipeline (#1636)

* fix rknn quantization pipeline

* resolve comments
This commit is contained in:
AllentDan 2023-01-31 16:35:11 +08:00 committed by GitHub
parent a19cf64765
commit 5ab0bcb5ea
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 3 additions and 7 deletions

View File

@ -50,14 +50,9 @@ class RKNNManager(BaseBackendManager):
"""
import importlib
try:
ret = importlib.util.find_spec('rknn-toolkit2') is not None
ret = importlib.util.find_spec('rknn') is not None
except Exception:
pass
if ret is None:
try:
ret = importlib.util.find_spec('rknn-toolkit') is not None
except Exception:
pass
return ret
@classmethod

View File

@ -218,7 +218,8 @@ def main():
dataset_file = tempfile.NamedTemporaryFile(suffix='.txt').name
with open(dataset_file, 'w') as f:
f.writelines([osp.abspath(args.img)])
quantization_cfg.setdefault('dataset', dataset_file)
if quantization_cfg.get('dataset', None) is None:
quantization_cfg['dataset'] = dataset_file
if backend == Backend.ASCEND:
# TODO: Add this to backend manager in the future
if args.dump_info: