Refactor onnx providers (#13804)

* Ensure dirs exist for model caches

* Formatting

* Don't use tensorrt for embeddings
This commit is contained in:
Nicolas Mowen 2024-09-17 14:54:44 -06:00 committed by GitHub
parent 38ff46e45c
commit 1ed8642010
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 21 additions and 9 deletions

View File

@ -1,4 +1,5 @@
import logging import logging
import os
import cv2 import cv2
import numpy as np import numpy as np
@ -42,6 +43,9 @@ class ONNXDetector(DetectionApi):
for provider in providers: for provider in providers:
if provider == "TensorrtExecutionProvider": if provider == "TensorrtExecutionProvider":
os.makedirs(
"/config/model_cache/tensorrt/ort/trt-engines", exist_ok=True
)
options.append( options.append(
{ {
"trt_timing_cache_enable": True, "trt_timing_cache_enable": True,
@ -51,7 +55,13 @@ class ONNXDetector(DetectionApi):
} }
) )
elif provider == "OpenVINOExecutionProvider": elif provider == "OpenVINOExecutionProvider":
options.append({"cache_dir": "/config/model_cache/openvino/ort"}) os.makedirs("/config/model_cache/openvino/ort", exist_ok=True)
options.append(
{
"cache_dir": "/config/model_cache/openvino/ort",
"device_type": "GPU",
}
)
else: else:
options.append({}) options.append({})

View File

@ -49,22 +49,24 @@ class Clip(OnnxClip):
@staticmethod @staticmethod
def _load_model(path: str, silent: bool): def _load_model(path: str, silent: bool):
providers = ort.get_available_providers() providers = []
options = [] options = []
for provider in providers: for provider in ort.get_available_providers():
if provider == "TensorrtExecutionProvider": if provider == "TensorrtExecutionProvider":
continue
elif provider == "OpenVINOExecutionProvider":
# TODO need to verify openvino works correctly
os.makedirs("/config/model_cache/openvino/ort", exist_ok=True)
providers.append(provider)
options.append( options.append(
{ {
"trt_timing_cache_enable": True, "cache_dir": "/config/model_cache/openvino/ort",
"trt_timing_cache_path": "/config/model_cache/tensorrt/ort", "device_type": "GPU",
"trt_engine_cache_enable": True,
"trt_engine_cache_path": "/config/model_cache/tensorrt/ort/trt-engines",
} }
) )
elif provider == "OpenVINOExecutionProvider":
options.append({"cache_dir": "/config/model_cache/openvino/ort"})
else: else:
providers.append(provider)
options.append({}) options.append({})
try: try: