d935d30e4c
Currently ONNXRuntime's CoreML support is not very good, and on my environment, I get errors like “coreml_execution_provider.cc:192 operator() Input (/sdp/flows.3/GatherND_2_output_0) has a dynamic shape ({-1,-1}) but the runtime shape ({0,10}) has zero elements. This is not supported by the CoreML EP.” and inference fails.
I hope this will be fixed in a future version of ONNXRuntime, and leave the test code as it is.
18 lines
712 B
Python
18 lines
712 B
Python
from typing import Any, Sequence, Union
|
|
|
|
|
|
def torch_device_to_onnx_providers(
|
|
device: str,
|
|
) -> Sequence[Union[str, tuple[str, dict[str, Any]]]]:
|
|
if device.startswith("cuda"):
|
|
return [
|
|
# cudnn_conv_algo_search を DEFAULT にすると推論速度が大幅に向上する
|
|
# ref: https://medium.com/neuml/debug-onnx-gpu-performance-c9290fe07459
|
|
("CUDAExecutionProvider", {"cudnn_conv_algo_search": "DEFAULT"}),
|
|
# CUDA が利用できない場合、可能であれば DirectML を利用する
|
|
("DmlExecutionProvider", {"device_id": 0}),
|
|
("CPUExecutionProvider", {}),
|
|
]
|
|
else:
|
|
return ["CPUExecutionProvider"]
|