diff --git a/machine-learning/rknn/export/build_rknn.py b/machine-learning/rknn/export/build_rknn.py new file mode 100644 index 0000000000..15c2da04a2 --- /dev/null +++ b/machine-learning/rknn/export/build_rknn.py @@ -0,0 +1,53 @@ +import argparse +import os + +parser = argparse.ArgumentParser("RKNN model converting") +parser.add_argument("model", help="Directory of the model that will be exported to RKNN ex:ViT-B-32__openai.", type=str) +parser.add_argument("target_platform", help="target platform ex:rk3566", type=str) +args = parser.parse_args() + + +def ConvertModel(model_path='ViT-B-32__openai/textual/model.onnx', target_platform='rk3566', dynamic_input = None): + # E build: Repeat call the 'rknn.build' or 'rknn.hybrid_quantization_step1' is not allow! + from rknn.api import RKNN + rknn = RKNN(verbose=False) + + rknn.config(target_platform=target_platform, dynamic_input=dynamic_input) + ret = rknn.load_onnx(model=model_path) + + if ret != 0: + print("Load failed!") + exit(ret) + + ret = rknn.build(do_quantization=False) + + if ret != 0: + print("Build failed!") + exit(ret) + print(model_path.replace('model.onnx',f'{target_platform}.rknn')) + ret = rknn.export_rknn(model_path.replace('model.onnx',f'{target_platform}.rknn')) + if ret != 0: + print('Export rknn model failed!') + exit(ret) + print('done') + del rknn + del RKNN + + if not os.path.isfile(f'{model_path.replace("onnx","rknn")}'): + print(f'Dummy model not found at {model_path.replace("onnx","rknn")}, creating one') + with open(f'{model_path.replace("onnx","rknn")}', 'w'): + pass + + +if os.path.isdir(f'{args.model}/textual') and os.path.isdir(f'{args.model}/visual'): # is a clip model + print('Converting Clip model.') + ConvertModel(model_path=f'{args.model}/textual/model.onnx', target_platform=args.target_platform) + ConvertModel(model_path=f'{args.model}/visual/model.onnx', target_platform=args.target_platform) + +elif os.path.isdir(f'{args.model}/detection') and os.path.isdir(f'{args.model}/recognition'): # is a facial model + print('Converting facial model.') + ConvertModel(f'{args.model}/detection/model.onnx', args.target_platform, [[[1, 3, 640, 640]]]) + ConvertModel(f'{args.model}/recognition/model.onnx', args.target_platform, [[[1, 3, 112, 112]]]) + +else: + print('Unknown model.') diff --git a/machine-learning/rknn/export/requirements.txt b/machine-learning/rknn/export/requirements.txt new file mode 100644 index 0000000000..982b7fcdb4 --- /dev/null +++ b/machine-learning/rknn/export/requirements.txt @@ -0,0 +1 @@ +rknn-toolkit2==2.3.0 \ No newline at end of file