Ensure dirs exist for model caches

This commit is contained in:
Nicolas Mowen 2024-09-17 14:39:27 -06:00
parent 2362d0e838
commit bb305c6d9d
2 changed files with 13 additions and 1 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

@ -54,6 +54,7 @@ class Clip(OnnxClip):
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,
@ -63,6 +64,7 @@ class Clip(OnnxClip):
} }
) )
elif provider == "OpenVINOExecutionProvider": elif provider == "OpenVINOExecutionProvider":
os.makedirs("/config/model_cache/openvino/ort", exist_ok=True)
options.append({"cache_dir": "/config/model_cache/openvino/ort"}) options.append({"cache_dir": "/config/model_cache/openvino/ort"})
else: else:
options.append({}) options.append({})