Don't use trt, use cuda

This commit is contained in:
Nicolas Mowen 2024-10-13 18:09:02 -06:00
parent 833768172d
commit 71cb9e7fd9

View File

@ -25,28 +25,23 @@ def get_ort_providers(
], ],
) )
providers = ort.get_available_providers() providers = []
options = [] options = []
for provider in providers: for provider in providers:
if provider == "TensorrtExecutionProvider": if provider == "CudaExecutionProvider":
os.makedirs("/config/model_cache/tensorrt/ort/trt-engines", exist_ok=True) providers.append(provider)
options.append(
if not requires_fp16 or os.environ.get("USE_FP_16", "True") != "False": {
options.append( "arena_extend_strategy": "kSameAsRequested",
{ }
"arena_extend_strategy": "kSameAsRequested", )
"trt_fp16_enable": requires_fp16, elif provider == "TensorrtExecutionProvider":
"trt_timing_cache_enable": True, # TensorrtExecutionProvider uses too much memory without options to control it
"trt_engine_cache_enable": True, pass
"trt_timing_cache_path": "/config/model_cache/tensorrt/ort",
"trt_engine_cache_path": "/config/model_cache/tensorrt/ort/trt-engines",
}
)
else:
options.append({})
elif provider == "OpenVINOExecutionProvider": elif provider == "OpenVINOExecutionProvider":
os.makedirs("/config/model_cache/openvino/ort", exist_ok=True) os.makedirs("/config/model_cache/openvino/ort", exist_ok=True)
providers.append(provider)
options.append( options.append(
{ {
"arena_extend_strategy": "kSameAsRequested", "arena_extend_strategy": "kSameAsRequested",
@ -55,12 +50,14 @@ def get_ort_providers(
} }
) )
elif provider == "CPUExecutionProvider": elif provider == "CPUExecutionProvider":
providers.append(provider)
options.append( options.append(
{ {
"arena_extend_strategy": "kSameAsRequested", "arena_extend_strategy": "kSameAsRequested",
} }
) )
else: else:
providers.append(provider)
options.append({}) options.append({})
return (providers, options) return (providers, options)