fix(device): закреплён ONNX CPU для автоматического выбора
This commit is contained in:
1 parent
2b54f5bc9e
commit
937eafe8db
4 files changed
+35
-15
No files matched your search
@@ -200,8 +200,9 @@ class OnnxAsrBackend:
|
|||||||
model = onnx_asr.load_model(
|
model = onnx_asr.load_model(
|
||||||
model=model_path,
|
model=model_path,
|
||||||
quantization=quantization,
|
quantization=quantization,
|
||||||
|
providers=["CPUExecutionProvider"],
|
||||||
)
|
)
|
||||||
vad = onnx_asr.load_vad("silero")
|
vad = onnx_asr.load_vad("silero", providers=["CPUExecutionProvider"])
|
||||||
self._vad = vad
|
self._vad = vad
|
||||||
return model.with_vad(vad).with_timestamps()
|
return model.with_vad(vad).with_timestamps()
|
||||||
|
|
||||||
@@ -278,7 +279,7 @@ class OnnxAsrBackend:
|
|||||||
fallback = _whisper_fallback_model(model_name)
|
fallback = _whisper_fallback_model(model_name)
|
||||||
raise ValueError(
|
raise ValueError(
|
||||||
f"Модель '{model_name}' относится к Whisper и не поддерживается "
|
f"Модель '{model_name}' относится к Whisper и не поддерживается "
|
||||||
"ONNX-бэкендом. Без CUDA --device auto выбирает ONNX; "
|
"ONNX-бэкендом. --device auto всегда выбирает ONNX CPU; "
|
||||||
f"укажите --device openvino-cpu --model {model_name} на x86, "
|
f"укажите --device openvino-cpu --model {model_name} на x86, "
|
||||||
f"--device cpu --model {fallback} на любой платформе "
|
f"--device cpu --model {fallback} на любой платформе "
|
||||||
f"или --device cuda --model {fallback} при NVIDIA GPU."
|
f"или --device cuda --model {fallback} при NVIDIA GPU."
|
||||||
|
|||||||
@@ -2,7 +2,6 @@
|
|||||||
|
|
||||||
import glob
|
import glob
|
||||||
import platform
|
import platform
|
||||||
import shutil
|
|
||||||
import subprocess
|
import subprocess
|
||||||
import warnings
|
import warnings
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
@@ -16,15 +15,13 @@ SUPPORTED_EXTENSIONS = {
|
|||||||
def detect_device(requested: str = "auto") -> str:
|
def detect_device(requested: str = "auto") -> str:
|
||||||
"""Определяет устройство для вычислений.
|
"""Определяет устройство для вычислений.
|
||||||
|
|
||||||
При ``requested="auto"`` выбирает CUDA при наличии ``nvidia-smi``,
|
При ``requested="auto"`` всегда выбирает ONNX на CPU.
|
||||||
иначе ONNX на CPU.
|
CUDA включается только явным выбором пользователя.
|
||||||
``"openvino"`` — алиас для авто-детекта внутри OpenVINO
|
``"openvino"`` — алиас для авто-детекта внутри OpenVINO
|
||||||
(GPU если доступен, иначе CPU).
|
(GPU если доступен, иначе CPU).
|
||||||
Возвращает всегда конкретное значение (не абстрактный ``"openvino"``).
|
Возвращает всегда конкретное значение (не абстрактный ``"openvino"``).
|
||||||
"""
|
"""
|
||||||
if requested == "auto":
|
if requested == "auto":
|
||||||
if shutil.which("nvidia-smi") is not None:
|
|
||||||
return "cuda"
|
|
||||||
return "onnx"
|
return "onnx"
|
||||||
if requested == "openvino":
|
if requested == "openvino":
|
||||||
if _is_openvino_gpu_available():
|
if _is_openvino_gpu_available():
|
||||||
|
|||||||
+25
-3
@@ -1,6 +1,7 @@
|
|||||||
"""Tests for onnx-asr backend."""
|
"""Tests for onnx-asr backend."""
|
||||||
|
|
||||||
import warnings
|
import warnings
|
||||||
|
from unittest.mock import MagicMock
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
|
|
||||||
@@ -64,12 +65,12 @@ class TestEnsureModelAvailable:
|
|||||||
def with_timestamps(self):
|
def with_timestamps(self):
|
||||||
return self
|
return self
|
||||||
|
|
||||||
def fake_load_model(*, model, quantization):
|
def fake_load_model(*, model, quantization, providers):
|
||||||
quantizations.append(quantization)
|
quantizations.append(quantization)
|
||||||
return FakeAsrAdapter()
|
return FakeAsrAdapter()
|
||||||
|
|
||||||
monkeypatch.setattr("onnx_asr.load_model", fake_load_model)
|
monkeypatch.setattr("onnx_asr.load_model", fake_load_model)
|
||||||
monkeypatch.setattr("onnx_asr.load_vad", lambda model: None)
|
monkeypatch.setattr("onnx_asr.load_vad", lambda model, **kwargs: None)
|
||||||
|
|
||||||
backend = OnnxAsrBackend(compute_type_explicit=False)
|
backend = OnnxAsrBackend(compute_type_explicit=False)
|
||||||
model_id = backend.ensure_model_available(
|
model_id = backend.ensure_model_available(
|
||||||
@@ -87,6 +88,26 @@ class TestEnsureModelAvailable:
|
|||||||
|
|
||||||
|
|
||||||
class TestCreateModel:
|
class TestCreateModel:
|
||||||
|
def test_cpu_provider_is_explicit_for_asr_and_vad(self, monkeypatch):
|
||||||
|
"""Доступность CoreML/CUDA не меняет исполнение ASR и VAD."""
|
||||||
|
monkeypatch.setattr(
|
||||||
|
"onnxruntime.get_available_providers",
|
||||||
|
lambda: ["CoreMLExecutionProvider", "CUDAExecutionProvider", "CPUExecutionProvider"],
|
||||||
|
)
|
||||||
|
load_asr = MagicMock()
|
||||||
|
load_vad = MagicMock()
|
||||||
|
monkeypatch.setattr("onnx_asr.load_model", load_asr)
|
||||||
|
monkeypatch.setattr("onnx_asr.load_vad", load_vad)
|
||||||
|
|
||||||
|
OnnxAsrBackend().create_model("gigaam-v3-e2e-rnnt", "onnx", "int8")
|
||||||
|
|
||||||
|
load_asr.assert_called_once_with(
|
||||||
|
model="gigaam-v3-e2e-rnnt",
|
||||||
|
quantization="int8",
|
||||||
|
providers=["CPUExecutionProvider"],
|
||||||
|
)
|
||||||
|
load_vad.assert_called_once_with("silero", providers=["CPUExecutionProvider"])
|
||||||
|
|
||||||
def test_wraps_vad_model_with_timestamps(self, monkeypatch):
|
def test_wraps_vad_model_with_timestamps(self, monkeypatch):
|
||||||
timestamped_model = object()
|
timestamped_model = object()
|
||||||
|
|
||||||
@@ -99,7 +120,7 @@ class TestCreateModel:
|
|||||||
return FakeVadAdapter()
|
return FakeVadAdapter()
|
||||||
|
|
||||||
monkeypatch.setattr("onnx_asr.load_model", lambda **kwargs: FakeAsrAdapter())
|
monkeypatch.setattr("onnx_asr.load_model", lambda **kwargs: FakeAsrAdapter())
|
||||||
monkeypatch.setattr("onnx_asr.load_vad", lambda model: object())
|
monkeypatch.setattr("onnx_asr.load_vad", lambda model, **kwargs: object())
|
||||||
|
|
||||||
model = OnnxAsrBackend().create_model("gigaam-v3-e2e-rnnt", "onnx", "int8")
|
model = OnnxAsrBackend().create_model("gigaam-v3-e2e-rnnt", "onnx", "int8")
|
||||||
|
|
||||||
@@ -127,6 +148,7 @@ class TestCreateModel:
|
|||||||
return self
|
return self
|
||||||
|
|
||||||
monkeypatch.setattr("onnx_asr.load_model", fake_load_model)
|
monkeypatch.setattr("onnx_asr.load_model", fake_load_model)
|
||||||
|
monkeypatch.setattr("onnx_asr.load_vad", lambda model, **kwargs: None)
|
||||||
|
|
||||||
backend = OnnxAsrBackend()
|
backend = OnnxAsrBackend()
|
||||||
backend.actual_compute_type = "int8"
|
backend.actual_compute_type = "int8"
|
||||||
|
|||||||
+5
-5
@@ -66,17 +66,17 @@ def test_detect_device_explicit_passthrough():
|
|||||||
assert detect_device("openvino-cpu") == "openvino-cpu"
|
assert detect_device("openvino-cpu") == "openvino-cpu"
|
||||||
|
|
||||||
|
|
||||||
def test_detect_device_auto_cuda_when_nvidia_smi_available():
|
def test_detect_device_auto_onnx_when_nvidia_smi_available():
|
||||||
"""При доступном nvidia-smi auto выбирает CUDA."""
|
"""При доступном nvidia-smi auto остаётся на ONNX CPU."""
|
||||||
with patch(
|
with patch(
|
||||||
"local_transcriber.utils.shutil.which", return_value="/usr/bin/nvidia-smi"
|
"shutil.which", return_value="/usr/bin/nvidia-smi"
|
||||||
):
|
):
|
||||||
assert detect_device("auto") == "cuda"
|
assert detect_device("auto") == "onnx"
|
||||||
|
|
||||||
|
|
||||||
def test_detect_device_auto_onnx_without_cuda():
|
def test_detect_device_auto_onnx_without_cuda():
|
||||||
"""Без CUDA auto выбирает ONNX CPU."""
|
"""Без CUDA auto выбирает ONNX CPU."""
|
||||||
with patch("local_transcriber.utils.shutil.which", return_value=None):
|
with patch("shutil.which", return_value=None):
|
||||||
assert detect_device("auto") == "onnx"
|
assert detect_device("auto") == "onnx"
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
Reference in new issue
Block a user