When pymilvus[model] opens an ONNX model, it doesn't set the execution provider, which will default to CPU: https://github.com/milvus-io/milvus-model/blob/main/src/pymilvus/model/dense/onnx.py#L20. Typically you would want to use your GPU to accelerate inference...
How about we add some logic before that line so that it will use CUDA when that's available and CoreML on Macs?
I think this would work:
session = onnxruntime.InferenceSession(model, providers=['CUDAExecutionProvider', 'CoreMLExecutionProvider', 'CPUExecutionProvider'])
but if not:
if torch.backends.mps.is_available():
providers = 'CoreMLExecutionProvider'
elif torch.cuda.is_available():
providers = 'CUDAExecutionProvider'
else:
providers = 'CPUExecutionProvider'
session = onnxruntime.InferenceSession(model, providers=providers)
When pymilvus[model] opens an ONNX model, it doesn't set the execution provider, which will default to CPU: https://github.com/milvus-io/milvus-model/blob/main/src/pymilvus/model/dense/onnx.py#L20. Typically you would want to use your GPU to accelerate inference...
How about we add some logic before that line so that it will use CUDA when that's available and CoreML on Macs?
I think this would work:
but if not: