diff --git a/frigate/detectors/detection_runners.py b/frigate/detectors/detection_runners.py index ee30f4fead..a998b449ba 100644 --- a/frigate/detectors/detection_runners.py +++ b/frigate/detectors/detection_runners.py @@ -25,25 +25,31 @@ def is_arm64_platform() -> bool: return machine in ("aarch64", "arm64", "armv8", "armv7l") -def get_ort_session_options( - is_complex_model: bool = False, -) -> ort.SessionOptions | None: +def get_ort_session_options(model_type: str | None = None) -> ort.SessionOptions | None: """Get ONNX Runtime session options with appropriate settings. Args: - is_complex_model: Whether the model needs basic optimization to avoid graph fusion issues. + model_type: Model being loaded, used to pin its graph optimization level. Returns: - SessionOptions with appropriate optimization level, or None for default settings. + SessionOptions with a pinned optimization level, or None for default settings. """ - if is_complex_model: - sess_options = ort.SessionOptions() - sess_options.graph_optimization_level = ( - ort.GraphOptimizationLevel.ORT_ENABLE_BASIC - ) - return sess_options + # Import here to avoid circular imports + from frigate.embeddings.types import EnrichmentModelTypeEnum - return None + if model_type == EnrichmentModelTypeEnum.jina_v2.value: + # below EXTENDED the CUDA EP returns an identical vector for every image, + # and ORT_ENABLE_ALL fails to build on CPU with a SimplifiedLayerNormFusion error + level = ort.GraphOptimizationLevel.ORT_ENABLE_EXTENDED + elif model_type == EnrichmentModelTypeEnum.jina_v1.value: + # aggressive optimizations create or expect nodes that don't exist + level = ort.GraphOptimizationLevel.ORT_ENABLE_BASIC + else: + return None + + sess_options = ort.SessionOptions() + sess_options.graph_optimization_level = level + return sess_options # Import OpenVINO only when needed to avoid circular dependencies @@ -115,21 +121,6 @@ class BaseModelRunner(ABC): class ONNXModelRunner(BaseModelRunner): """Run ONNX models using ONNX Runtime.""" - @staticmethod - def is_cpu_complex_model(model_type: str) -> bool: - """Check if model needs basic optimization level to avoid graph fusion issues. - - Some models (like Jina-CLIP) have issues with aggressive optimizations like - SimplifiedLayerNormFusion that create or expect nodes that don't exist. - """ - # Import here to avoid circular imports - from frigate.embeddings.types import EnrichmentModelTypeEnum - - return model_type in [ - EnrichmentModelTypeEnum.jina_v1.value, - EnrichmentModelTypeEnum.jina_v2.value, - ] - @staticmethod def is_migraphx_complex_model(model_type: str) -> bool: # Import here to avoid circular imports @@ -626,9 +617,7 @@ def get_optimized_runner( return ONNXModelRunner( ort.InferenceSession( model_path, - sess_options=get_ort_session_options( - ONNXModelRunner.is_cpu_complex_model(model_type) - ), + sess_options=get_ort_session_options(model_type), providers=providers, provider_options=options, ), diff --git a/frigate/test/test_detection_runners.py b/frigate/test/test_detection_runners.py new file mode 100644 index 0000000000..fbc71161d4 --- /dev/null +++ b/frigate/test/test_detection_runners.py @@ -0,0 +1,41 @@ +"""Tests for ONNX Runtime session option selection.""" + +import unittest + +import onnxruntime as ort + +from frigate.detectors.detection_runners import get_ort_session_options +from frigate.detectors.detector_config import ModelTypeEnum +from frigate.embeddings.types import EnrichmentModelTypeEnum + + +class TestGetOrtSessionOptions(unittest.TestCase): + def test_jina_v2_uses_extended(self): + """jina-clip-v2 returns an identical vector for every image on the CUDA + execution provider at anything below EXTENDED.""" + options = get_ort_session_options(EnrichmentModelTypeEnum.jina_v2.value) + + self.assertIsNotNone(options) + self.assertEqual( + options.graph_optimization_level, + ort.GraphOptimizationLevel.ORT_ENABLE_EXTENDED, + ) + + def test_jina_v1_uses_basic(self): + options = get_ort_session_options(EnrichmentModelTypeEnum.jina_v1.value) + + self.assertIsNotNone(options) + self.assertEqual( + options.graph_optimization_level, + ort.GraphOptimizationLevel.ORT_ENABLE_BASIC, + ) + + def test_other_models_use_defaults(self): + for model_type in [ + None, + EnrichmentModelTypeEnum.paddleocr.value, + EnrichmentModelTypeEnum.arcface.value, + ModelTypeEnum.rfdetr.value, + ]: + with self.subTest(model_type=model_type): + self.assertIsNone(get_ort_session_options(model_type))