mirrors
/
PaddleOCR
mirror of https://github.com/PaddlePaddle/PaddleOCR.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

solve the onnxruntime inference issue (#13154)

pull/13105/head^2
CaiRan 2024-06-21 17:20:40 +08:00 committed by GitHub
parent f8ca01dc01
commit 6d5f998fe1
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 10 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
tools/infer/utility.py
View File

@ -197,10 +197,18 @@ def create_predictor(args, mode, logger):
raise ValueError("not find model file path {}".format(model_file_path)) raise ValueError("not find model file path {}".format(model_file_path))
if args.use_gpu: if args.use_gpu:
sess = ort.InferenceSession( sess = ort.InferenceSession(
model_file_path, providers=["CUDAExecutionProvider"] model_file_path,
providers=[
(
"CUDAExecutionProvider",
{"device_id": args.gpu_id, "cudnn_conv_algo_search": "DEFAULT"},
)
],
) )
else: else:
sess = ort.InferenceSession(model_file_path) sess = ort.InferenceSession(
model_file_path, providers=["CPUExecutionProvider"]
)
return sess, sess.get_inputs()[0], None, None return sess, sess.get_inputs()[0], None, None
else: else: