使用gpu调用文档检测模型

2024-08-27 17:11:19 +08:00
parent 4009d7eb81
commit ec0e2082c0
2 changed files with 2 additions and 2 deletions
--- a/paddle_detection/detector.py
+++ b/paddle_detection/detector.py
@@ -33,7 +33,7 @@ def detect_image(img_path):
    onnx_file = "model/object_det_model/ppyoloe_plus_crn_x_80e_coco_w_nms.onnx"
    infer_cfg = "model/object_det_model/infer_cfg.yml"
    # load predictor
-    predictor = InferenceSession(onnx_file)
+    predictor = InferenceSession(onnx_file, providers=["CUDAExecutionProvider"])
    # load infer config
    infer_config = PredictConfig(infer_cfg)