self-hosted/immich
1
0
Fork 0
mirror of https://github.com/immich-app/immich.git synced 2025-08-08 23:07:06 +02:00
Code Issues Releases Activity

feat(ml): round-robin device assignment (#13237)

Browse Source 
* round-robin device assignment

* docs and tests

clarify doc
This commit is contained in:
Mert
2024-10-07 17:37:45 -04:00
committed by GitHub
parent 063969ca05
commit bd826b0b9b
8 changed files with 62 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
machine-learning/app/sessions/ort.py
View File

@ -86,11 +86,13 @@ class OrtSession:
provider_options = []
for provider in self.providers:
match provider:
case "CPUExecutionProvider" | "CUDAExecutionProvider":
case "CPUExecutionProvider":
options = {"arena_extend_strategy": "kSameAsRequested"}
case "CUDAExecutionProvider":
options = {"arena_extend_strategy": "kSameAsRequested", "device_id": settings.device_id}
case "OpenVINOExecutionProvider":
options = {
"device_type": "GPU",
"device_type": f"GPU.{settings.device_id}",
"precision": "FP32",
"cache_dir": (self.model_path.parent / "openvino").as_posix(),
}