添加{"cudnn_conv_algo_search": "DEFAULT"}参数，加速gpu下onnxruntime推理速度；指定"CP…

…UExecutionProvider"参数解决cpu环境下onnxruntime报错
PaddlePaddle · Jun 20, 2024 · 0f63d56 · 0f63d56
1 parent f8ca01d
commit 0f63d56
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tools/infer/utility.py b/tools/infer/utility.py
@@ -197,10 +197,10 @@ def create_predictor(args, mode, logger):
             raise ValueError("not find model file path {}".format(model_file_path))
         if args.use_gpu:
             sess = ort.InferenceSession(
-                model_file_path, providers=["CUDAExecutionProvider"]
+                model_file_path, providers=[("CUDAExecutionProvider", {"cudnn_conv_algo_search": "DEFAULT"})]
             )
         else:
-            sess = ort.InferenceSession(model_file_path)
+            sess = ort.InferenceSession(model_file_path, providers=["CPUExecutionProvider"])
         return sess, sess.get_inputs()[0], None, None
 
     else: