chris/frigate
1
0
Fork 0
forked from Github/frigate
Code Pull Requests Packages Activity

Support ONNX model caching (#13780)

Browse Source 
* Support model caching

* Cleanup
This commit is contained in:
Nicolas Mowen
2024-09-16 18:18:11 -06:00
committed by GitHub
parent 4fc8d33d31
commit 36d7eb7caa
3 changed files with 27 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

23
frigate/detectors/plugins/onnx.py
View File

@@ -36,7 +36,28 @@ class ONNXDetector(DetectionApi):
path = detector_config.model.path
logger.info(f"ONNX: loading {detector_config.model.path}")
self.model = ort.InferenceSession(path, providers=ort.get_available_providers())
providers = ort.get_available_providers()
options = []
for provider in providers:
if provider == "TensorrtExecutionProvider":
options.append(
{
"trt_timing_cache_enable": True,
"trt_timing_cache_path": "/config/model_cache/tensorrt/ort",
"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:
options.append({})
self.model = ort.InferenceSession(
path, providers=providers, provider_options=options
)
self.h = detector_config.model.height
self.w = detector_config.model.width