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 os
import cv2
import numpy as np
@ -42,6 +43,9 @@ class ONNXDetector(DetectionApi):
for provider in providers:
if provider == "TensorrtExecutionProvider":
os.makedirs(
"/config/model_cache/tensorrt/ort/trt-engines", exist_ok=True
)
options.append(
{
"trt_timing_cache_enable": True,
@ -51,7 +55,13 @@ class ONNXDetector(DetectionApi):
}
)
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:
options.append({})

View File

@ -54,6 +54,7 @@ class Clip(OnnxClip):
for provider in providers:
if provider == "TensorrtExecutionProvider":
os.makedirs("/config/model_cache/tensorrt/ort/trt-engines", exist_ok=True)
options.append(
{
"trt_timing_cache_enable": True,
@ -63,6 +64,7 @@ class Clip(OnnxClip):
}
)
elif provider == "OpenVINOExecutionProvider":
os.makedirs("/config/model_cache/openvino/ort", exist_ok=True)
options.append({"cache_dir": "/config/model_cache/openvino/ort"})
else:
options.append({})