Spaces:
Paused
Paused
File size: 1,541 Bytes
e04e4d5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 |
from typing import Any, List
import onnxruntime
def encode_execution_providers(execution_providers : List[str]) -> List[str]:
return [ execution_provider.replace('ExecutionProvider', '').lower() for execution_provider in execution_providers ]
def decode_execution_providers(execution_providers: List[str]) -> List[str]:
available_execution_providers = onnxruntime.get_available_providers()
encoded_execution_providers = encode_execution_providers(available_execution_providers)
return [ execution_provider for execution_provider, encoded_execution_provider in zip(available_execution_providers, encoded_execution_providers) if any(execution_provider in encoded_execution_provider for execution_provider in execution_providers) ]
def apply_execution_provider_options(execution_providers: List[str]) -> List[Any]:
execution_providers_with_options : List[Any] = []
for execution_provider in execution_providers:
if execution_provider == 'CUDAExecutionProvider':
execution_providers_with_options.append((execution_provider,
{
'cudnn_conv_algo_search': 'DEFAULT'
}))
else:
execution_providers_with_options.append(execution_provider)
return execution_providers_with_options
def map_torch_backend(execution_providers : List[str]) -> str:
if 'CoreMLExecutionProvider' in execution_providers:
return 'mps'
if 'CUDAExecutionProvider' in execution_providers or 'ROCMExecutionProvider' in execution_providers :
return 'cuda'
if 'OpenVINOExecutionProvider' in execution_providers:
return 'mkl'
return 'cpu'
|