feat: GPU acceleration support for ONNX Runtime inference

Centralize execution provider selection in providers.py with
auto-detection and graceful fallback. All ONNX sessions (embedder,
detector, face processor, recognizer) now use the configured providers.

- New VISION_EXECUTION_PROVIDERS env var: "auto" for GPU auto-detect,
  or explicit "CUDAExecutionProvider,CPUExecutionProvider"
- Provider priority: CUDA > ROCm > OpenVINO > CPU (when set to "auto")
- docker-compose.yml includes commented-out NVIDIA GPU deploy section
- Supports onnxruntime-gpu as a drop-in replacement for onnxruntime

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
2026-04-12 22:12:21 +02:00
parent 94c07b1d0d
commit bbb8e4850c
7 changed files with 138 additions and 20 deletions

View File

@@ -66,6 +66,11 @@ class VisionSettings(BaseModel):
enabled: bool = True enabled: bool = True
backend: str = "onnx" # "onnx" | "rocm" (future) backend: str = "onnx" # "onnx" | "rocm" (future)
models_dir: str = "/data/models" models_dir: str = "/data/models"
# ONNX Runtime execution providers in priority order.
# Auto-detected at startup; falls back to CPU if GPU is unavailable.
# Options: "CUDAExecutionProvider", "ROCMExecutionProvider",
# "OpenVINOExecutionProvider", "CPUExecutionProvider"
execution_providers: list[str] = ["CPUExecutionProvider"]
embedder: EmbedderSettings = EmbedderSettings() embedder: EmbedderSettings = EmbedderSettings()
ocr: OCRSettings = OCRSettings() ocr: OCRSettings = OCRSettings()
detector: DetectorSettings = DetectorSettings() detector: DetectorSettings = DetectorSettings()
@@ -183,9 +188,22 @@ class Settings(BaseSettings):
def performance(self) -> PerformanceSettings: def performance(self) -> PerformanceSettings:
return self.config.performance return self.config.performance
# ONNX Runtime execution providers, overridable via env var.
# Comma-separated: "CUDAExecutionProvider,CPUExecutionProvider"
# or "auto" for GPU auto-detection.
vision_execution_providers: str = Field(
default="CPUExecutionProvider",
env="VISION_EXECUTION_PROVIDERS",
)
@property @property
def vision(self) -> VisionSettings: def vision(self) -> VisionSettings:
return self.config.vision v = self.config.vision
# Override execution_providers from env if set.
providers = [p.strip() for p in self.vision_execution_providers.split(",") if p.strip()]
if providers:
v.execution_providers = providers
return v
class Config: class Config:
env_file = ".env" env_file = ".env"

View File

@@ -120,12 +120,10 @@ class YOLOv8Detector(ObjectDetector):
def __init__(self, settings: VisionSettings): def __init__(self, settings: VisionSettings):
model_path = Path(settings.models_dir) / "detect" / "yolov8n.onnx" model_path = Path(settings.models_dir) / "detect" / "yolov8n.onnx"
opts = ort.SessionOptions() from app.services.vision.providers import create_session
opts.inter_op_num_threads = 2
opts.intra_op_num_threads = 2
logger.info("Loading YOLOv8n from %s", model_path) logger.info("Loading YOLOv8n from %s", model_path)
self._session = ort.InferenceSession(str(model_path), opts, providers=["CPUExecutionProvider"]) self._session = create_session(str(model_path), configured_providers=settings.execution_providers)
self._conf_threshold = settings.detector.min_confidence self._conf_threshold = settings.detector.min_confidence
self._max_detections = settings.detector.max_detections self._max_detections = settings.detector.max_detections

View File

@@ -67,15 +67,15 @@ class OpenCLIPEmbedder(Embedder):
visual_path = model_dir / "visual.onnx" visual_path = model_dir / "visual.onnx"
textual_path = model_dir / "textual.onnx" textual_path = model_dir / "textual.onnx"
opts = ort.SessionOptions() from app.services.vision.providers import create_session
opts.inter_op_num_threads = 2 from app.config import settings as app_settings
opts.intra_op_num_threads = 2 providers = app_settings.vision.execution_providers
logger.info("Loading OpenCLIP visual encoder from %s", visual_path) logger.info("Loading OpenCLIP visual encoder from %s", visual_path)
self._visual = ort.InferenceSession(str(visual_path), opts, providers=["CPUExecutionProvider"]) self._visual = create_session(str(visual_path), configured_providers=providers)
logger.info("Loading OpenCLIP textual encoder from %s", textual_path) logger.info("Loading OpenCLIP textual encoder from %s", textual_path)
self._textual = ort.InferenceSession(str(textual_path), opts, providers=["CPUExecutionProvider"]) self._textual = create_session(str(textual_path), configured_providers=providers)
def embed_image(self, image: np.ndarray) -> np.ndarray: def embed_image(self, image: np.ndarray) -> np.ndarray:
inp = _preprocess_image(image, _OPENCLIP_SIZE, _OPENCLIP_MEAN, _OPENCLIP_STD) inp = _preprocess_image(image, _OPENCLIP_SIZE, _OPENCLIP_MEAN, _OPENCLIP_STD)
@@ -110,15 +110,15 @@ class SigLIP2Embedder(Embedder):
visual_path = model_dir / "visual.onnx" visual_path = model_dir / "visual.onnx"
textual_path = model_dir / "textual.onnx" textual_path = model_dir / "textual.onnx"
opts = ort.SessionOptions() from app.services.vision.providers import create_session
opts.inter_op_num_threads = 2 from app.config import settings as app_settings
opts.intra_op_num_threads = 2 providers = app_settings.vision.execution_providers
logger.info("Loading SigLIP2 visual encoder from %s", visual_path) logger.info("Loading SigLIP2 visual encoder from %s", visual_path)
self._visual = ort.InferenceSession(str(visual_path), opts, providers=["CPUExecutionProvider"]) self._visual = create_session(str(visual_path), configured_providers=providers)
logger.info("Loading SigLIP2 textual encoder from %s", textual_path) logger.info("Loading SigLIP2 textual encoder from %s", textual_path)
self._textual = ort.InferenceSession(str(textual_path), opts, providers=["CPUExecutionProvider"]) self._textual = create_session(str(textual_path), configured_providers=providers)
def embed_image(self, image: np.ndarray) -> np.ndarray: def embed_image(self, image: np.ndarray) -> np.ndarray:
inp = _preprocess_image(image, _SIGLIP2_SIZE, _SIGLIP2_MEAN, _SIGLIP2_STD) inp = _preprocess_image(image, _SIGLIP2_SIZE, _SIGLIP2_MEAN, _SIGLIP2_STD)

View File

@@ -71,11 +71,9 @@ class YuNetSFaceProcessor(FaceProcessor):
logger.info("YuNet face detector loaded via OpenCV") logger.info("YuNet face detector loaded via OpenCV")
# SFace via ONNX Runtime # SFace via ONNX Runtime
opts = ort.SessionOptions() from app.services.vision.providers import create_session
opts.inter_op_num_threads = 2
opts.intra_op_num_threads = 2
ort.set_default_logger_severity(3) ort.set_default_logger_severity(3)
self._recognizer = ort.InferenceSession(sface_path, opts, providers=["CPUExecutionProvider"]) self._recognizer = create_session(sface_path, configured_providers=settings.execution_providers)
logger.info("SFace recognizer loaded via ONNX Runtime") logger.info("SFace recognizer loaded via ONNX Runtime")
self._min_face_size = settings.faces.min_face_size self._min_face_size = settings.faces.min_face_size

View File

@@ -23,10 +23,13 @@ class InsightFaceProcessor(FaceProcessor):
model_root = str(Path(settings.models_dir) / "face" / "insightface") model_root = str(Path(settings.models_dir) / "face" / "insightface")
logger.info("Loading InsightFace buffalo_l from %s", model_root) logger.info("Loading InsightFace buffalo_l from %s", model_root)
from app.services.vision.providers import get_providers
providers = get_providers(settings.execution_providers)
self._app = FaceAnalysis( self._app = FaceAnalysis(
name="buffalo_l", name="buffalo_l",
root=model_root, root=model_root,
providers=["CPUExecutionProvider"], providers=providers,
) )
self._app.prepare(ctx_id=-1, det_size=(640, 640)) self._app.prepare(ctx_id=-1, det_size=(640, 640))
self._min_det_score = settings.faces.recognition_threshold self._min_det_score = settings.faces.recognition_threshold

View File

@@ -0,0 +1,86 @@
"""
ONNX Runtime execution provider resolution with GPU auto-detection.
Resolves configured execution providers against what's actually available
in the current ONNX Runtime build. Falls back to CPU if no GPU provider
is available. Logs the selected provider so users can confirm GPU is active.
"""
import logging
import onnxruntime as ort
logger = logging.getLogger(__name__)
_resolved: list[str] | None = None
def get_providers(configured: list[str] | None = None) -> list[str]:
"""Return the best available execution providers.
1. If `configured` is provided, filter to only those that are
actually available in the current ORT build.
2. If none of the configured providers are available, fall back
to CPUExecutionProvider.
3. Auto-detect: if configured is ["auto"], probe for GPU providers.
Results are cached after first call.
"""
global _resolved
if _resolved is not None:
return _resolved
available = set(ort.get_available_providers())
logger.info("ONNX Runtime available providers: %s", sorted(available))
if configured is None or configured == ["CPUExecutionProvider"]:
_resolved = ["CPUExecutionProvider"]
return _resolved
if configured == ["auto"]:
# Auto-detect: prefer CUDA > ROCm > OpenVINO > CPU
priority = [
"CUDAExecutionProvider",
"ROCMExecutionProvider",
"OpenVINOExecutionProvider",
]
for p in priority:
if p in available:
_resolved = [p, "CPUExecutionProvider"]
logger.info("Auto-detected GPU provider: %s", p)
return _resolved
_resolved = ["CPUExecutionProvider"]
logger.info("No GPU provider detected, using CPU")
return _resolved
# Filter configured list to available providers.
resolved = [p for p in configured if p in available]
if not resolved:
logger.warning(
"None of the configured providers %s are available. "
"Falling back to CPU. Available: %s",
configured,
sorted(available),
)
resolved = ["CPUExecutionProvider"]
else:
# Always include CPU as fallback.
if "CPUExecutionProvider" not in resolved:
resolved.append("CPUExecutionProvider")
_resolved = resolved
logger.info("Using ONNX Runtime providers: %s", _resolved)
return _resolved
def create_session(
model_path: str,
opts: ort.SessionOptions | None = None,
configured_providers: list[str] | None = None,
) -> ort.InferenceSession:
"""Create an ONNX InferenceSession with the best available providers."""
providers = get_providers(configured_providers)
if opts is None:
opts = ort.SessionOptions()
opts.inter_op_num_threads = 2
opts.intra_op_num_threads = 2
return ort.InferenceSession(model_path, opts, providers=providers)

View File

@@ -143,6 +143,13 @@ services:
- LOG_LEVEL=${LOG_LEVEL:-INFO} - LOG_LEVEL=${LOG_LEVEL:-INFO}
- TZ=${TZ:-UTC} - TZ=${TZ:-UTC}
- MULITA_CELERY_WORKER=1 - MULITA_CELERY_WORKER=1
# ONNX Runtime execution providers. Set to "auto" to auto-detect
# GPU (CUDA > ROCm > OpenVINO > CPU), or explicitly:
# "CUDAExecutionProvider,CPUExecutionProvider"
# "ROCMExecutionProvider,CPUExecutionProvider"
# Default: CPU only. To enable GPU, also uncomment the deploy
# section below and install nvidia-container-toolkit on the host.
- VISION_EXECUTION_PROVIDERS=${VISION_EXECUTION_PROVIDERS:-CPUExecutionProvider}
# Pin each ONNX session to one intra-op thread so N prefork children # Pin each ONNX session to one intra-op thread so N prefork children
# × default-all-cores doesn't oversubscribe the box. With # × default-all-cores doesn't oversubscribe the box. With
# concurrency=5 and OMP=1, vision peaks at 5 busy cores, leaving # concurrency=5 and OMP=1, vision peaks at 5 busy cores, leaving
@@ -151,6 +158,14 @@ services:
- OMP_NUM_THREADS=1 - OMP_NUM_THREADS=1
- OPENBLAS_NUM_THREADS=1 - OPENBLAS_NUM_THREADS=1
- MKL_NUM_THREADS=1 - MKL_NUM_THREADS=1
# Uncomment for NVIDIA GPU passthrough:
# deploy:
# resources:
# reservations:
# devices:
# - driver: nvidia
# count: all
# capabilities: [gpu]
depends_on: depends_on:
redis: redis:
condition: service_started condition: service_started