Add ONNX inference providers (#5918)
* Add ONNX inference providers Fix for https://github.com/ultralytics/yolov5/issues/5916 * Update common.py
This commit is contained in:
1 parent
554f782537
commit
581dc301a7
1 file changed
+4
-2
+4
-2
@@ -320,9 +320,11 @@ class DetectMultiBackend(nn.Module):
|
||||
net = cv2.dnn.readNetFromONNX(w)
|
||||
elif onnx: # ONNX Runtime
|
||||
LOGGER.info(f'Loading {w} for ONNX Runtime inference...')
|
||||
check_requirements(('onnx', 'onnxruntime-gpu' if torch.cuda.is_available() else 'onnxruntime'))
|
||||
cuda = torch.cuda.is_available()
|
||||
check_requirements(('onnx', 'onnxruntime-gpu' if cuda else 'onnxruntime'))
|
||||
import onnxruntime
|
||||
session = onnxruntime.InferenceSession(w, None)
|
||||
providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] if cuda else ['CPUExecutionProvider']
|
||||
session = onnxruntime.InferenceSession(w, providers=providers)
|
||||
elif engine: # TensorRT
|
||||
LOGGER.info(f'Loading {w} for TensorRT inference...')
|
||||
import tensorrt as trt # https://developer.nvidia.com/nvidia-tensorrt-download
|
||||
|
||||
Reference in new issue
Block a user