Allow `create_predictor` function to accept array of ONNX Execution Providers ()

* pass onnx execution providers to create_predictor function

* added ability to provide onnxruntime SessionOptions

* added argument parser for onnx_sess_options

---------

Co-authored-by: ggolda <gleb.golda@streamlabs.com>
pull/14108/head
Salmondx 2024-10-25 06:21:39 -06:00 committed by GitHub
parent f7439b4db0
commit 37e177595b
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 16 additions and 2 deletions
tools/infer

View File

@ -147,6 +147,8 @@ def init_args():
parser.add_argument("--show_log", type=str2bool, default=True)
parser.add_argument("--use_onnx", type=str2bool, default=False)
parser.add_argument("--onnx_providers", nargs="+", type=str, default=False)
parser.add_argument("--onnx_sess_options", type=list, default=False)
# extended function
parser.add_argument(
@ -193,7 +195,16 @@ def create_predictor(args, mode, logger):
model_file_path = model_dir
if not os.path.exists(model_file_path):
raise ValueError("not find model file path {}".format(model_file_path))
if args.use_gpu:
sess_options = args.onnx_sess_options or None
if args.onnx_providers and len(args.onnx_providers) > 0:
sess = ort.InferenceSession(
model_file_path,
providers=args.onnx_providers,
sess_options=sess_options,
)
elif args.use_gpu:
sess = ort.InferenceSession(
model_file_path,
providers=[
@ -202,10 +213,13 @@ def create_predictor(args, mode, logger):
{"device_id": args.gpu_id, "cudnn_conv_algo_search": "DEFAULT"},
)
],
sess_options=sess_options,
)
else:
sess = ort.InferenceSession(
model_file_path, providers=["CPUExecutionProvider"]
model_file_path,
providers=["CPUExecutionProvider"],
sess_options=sess_options,
)
return sess, sess.get_inputs()[0], None, None