220 lines
9.6 KiB
Python
220 lines
9.6 KiB
Python
"""Diagnose the PyTorch / CUDA situation so the UI can explain *why* it's on CPU.
|
|
|
|
Pure (no Qt). ``gather()`` is the heavy part — it imports torch and shells out to
|
|
``nvidia-smi`` — so call it off the GUI thread. ``analyze()`` is fast formatting on
|
|
the gathered dict and figures out the most likely cause + concrete fix.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import re
|
|
import shutil
|
|
import subprocess
|
|
import sys
|
|
|
|
# pip indexes for the CUDA builds (see README). cu121 needs a driver with CUDA >= 12.1.
|
|
CUDA_WHEELS = {
|
|
"cu121": "https://download.pytorch.org/whl/cu121",
|
|
"cu118": "https://download.pytorch.org/whl/cu118",
|
|
}
|
|
|
|
_NO_WINDOW = 0x08000000 if sys.platform == "win32" else 0 # CREATE_NO_WINDOW
|
|
|
|
|
|
def _run_nvidia_smi() -> dict:
|
|
"""Probe the NVIDIA driver/GPU via nvidia-smi. Never raises."""
|
|
out: dict = {"found": False, "gpus": [], "driver_version": None, "cuda_driver": None}
|
|
exe = shutil.which("nvidia-smi")
|
|
if not exe and sys.platform == "win32":
|
|
candidate = r"C:\Windows\System32\nvidia-smi.exe"
|
|
exe = candidate if shutil.os.path.isfile(candidate) else None
|
|
if not exe:
|
|
return out
|
|
# GPU names + driver version (robust CSV form).
|
|
try:
|
|
r = subprocess.run(
|
|
[exe, "--query-gpu=name,driver_version", "--format=csv,noheader,nounits"],
|
|
capture_output=True, text=True, timeout=10, creationflags=_NO_WINDOW,
|
|
)
|
|
if r.returncode == 0:
|
|
out["found"] = True
|
|
for line in r.stdout.strip().splitlines():
|
|
parts = [p.strip() for p in line.split(",")]
|
|
if parts and parts[0]:
|
|
out["gpus"].append(parts[0])
|
|
if len(parts) > 1 and parts[1]:
|
|
out["driver_version"] = parts[1]
|
|
except Exception:
|
|
return out
|
|
# Max CUDA version the driver supports (only in the plain header).
|
|
try:
|
|
r2 = subprocess.run(
|
|
[exe], capture_output=True, text=True, timeout=10, creationflags=_NO_WINDOW,
|
|
)
|
|
m = re.search(r"CUDA Version:\s*([\d.]+)", r2.stdout)
|
|
if m:
|
|
out["cuda_driver"] = m.group(1)
|
|
except Exception:
|
|
pass
|
|
return out
|
|
|
|
|
|
def gather() -> dict:
|
|
"""Collect torch + NVIDIA facts. Never raises — missing torch/GPU are valid."""
|
|
info: dict = {
|
|
"installed": False,
|
|
"version": None, # torch.__version__ (e.g. "2.12.0+cpu")
|
|
"built_cuda": None, # torch.version.cuda (None on a CPU-only build)
|
|
"cuda_available": False,
|
|
"device_name": None, # the active GPU's name, if any
|
|
"import_error": None,
|
|
# filled by nvidia-smi:
|
|
"nvidia_smi": False,
|
|
"gpus": [],
|
|
"driver_version": None,
|
|
"cuda_driver": None,
|
|
}
|
|
try:
|
|
import torch
|
|
except Exception as exc:
|
|
info["import_error"] = str(exc)
|
|
else:
|
|
info["installed"] = True
|
|
info["version"] = getattr(torch, "__version__", None)
|
|
try:
|
|
info["built_cuda"] = torch.version.cuda
|
|
except Exception:
|
|
info["built_cuda"] = None
|
|
try:
|
|
info["cuda_available"] = bool(torch.cuda.is_available())
|
|
except Exception:
|
|
info["cuda_available"] = False
|
|
if info["cuda_available"]:
|
|
try:
|
|
info["device_name"] = torch.cuda.get_device_name(0)
|
|
except Exception:
|
|
info["device_name"] = None
|
|
|
|
smi = _run_nvidia_smi()
|
|
info["nvidia_smi"] = smi["found"]
|
|
info["gpus"] = smi["gpus"]
|
|
info["driver_version"] = smi["driver_version"]
|
|
info["cuda_driver"] = smi["cuda_driver"]
|
|
return info
|
|
|
|
|
|
def _ver_tuple(v: str | None) -> tuple[int, ...]:
|
|
try:
|
|
return tuple(int(x) for x in str(v).split(".")[:2])
|
|
except (ValueError, AttributeError):
|
|
return ()
|
|
|
|
|
|
def recommend_channel(info: dict) -> str:
|
|
"""Pick the pip CUDA wheel index that matches the driver (cu118 for older)."""
|
|
cd = _ver_tuple(info.get("cuda_driver"))
|
|
if cd and cd < (12, 1):
|
|
return "cu118"
|
|
return "cu121"
|
|
|
|
|
|
def install_command(channel: str = "cu121") -> str:
|
|
"""The pip commands to (re)install the chosen CUDA build.
|
|
|
|
Targets the **running interpreter** (``sys.executable -m pip``) so the command
|
|
hits the same venv that runs the app — not whatever ``pip`` is on PATH. (A common
|
|
trap: running bare ``pip`` in a global shell while torch lives in the project venv.)
|
|
"""
|
|
index = CUDA_WHEELS.get(channel, CUDA_WHEELS["cu121"])
|
|
py = sys.executable or "python"
|
|
q = f'"{py}"' if " " in py else py
|
|
return (
|
|
f"{q} -m pip uninstall -y torch torchvision\n"
|
|
f"{q} -m pip install torch torchvision --index-url {index}"
|
|
)
|
|
|
|
|
|
def analyze(info: dict) -> dict:
|
|
"""Turn the raw facts into {summary, details[list], steps, command}."""
|
|
installed = info.get("installed")
|
|
cuda = info.get("cuda_available")
|
|
built = info.get("built_cuda")
|
|
gpus = info.get("gpus") or []
|
|
smi = info.get("nvidia_smi")
|
|
driver = info.get("driver_version")
|
|
cuda_driver = info.get("cuda_driver")
|
|
channel = recommend_channel(info)
|
|
command = install_command(channel)
|
|
|
|
if not installed:
|
|
build_str = "не установлен"
|
|
elif built:
|
|
build_str = f"CUDA {built}"
|
|
else:
|
|
build_str = "CPU-only (+cpu)"
|
|
gpu_str = (
|
|
", ".join(gpus) if gpus
|
|
else ("не обнаружена" if smi or info.get("nvidia_smi") is False else "nvidia-smi не найден")
|
|
)
|
|
if not gpus and not smi:
|
|
gpu_str = "nvidia-smi не найден (нет драйвера NVIDIA?)"
|
|
details = [
|
|
f"PyTorch: {info.get('version') or 'не установлен'}",
|
|
f"Сборка PyTorch: {build_str}",
|
|
f"CUDA доступна в PyTorch: {'да' if cuda else 'нет'}",
|
|
f"Видеокарта (nvidia-smi): {gpu_str}",
|
|
f"Драйвер NVIDIA: {driver or '—'}",
|
|
f"Макс. CUDA драйвера: {cuda_driver or '—'}",
|
|
f"Интерпретатор (venv): {sys.executable}",
|
|
]
|
|
if info.get("import_error"):
|
|
details.append(f"Ошибка импорта torch: {info['import_error']}")
|
|
|
|
if not installed:
|
|
summary = "PyTorch не установлен — детекция YOLO и DeepMosaics идут на CPU."
|
|
steps = ("Установите PyTorch (CUDA-сборку, если есть NVIDIA-видеокарта):\n\n"
|
|
+ command + "\n\nДля YOLO также: pip install -e \".[yolo]\"")
|
|
elif cuda:
|
|
gpu = info.get("device_name") or (gpus[0] if gpus else "GPU")
|
|
summary = f"Всё в порядке: PyTorch использует CUDA. Активный GPU: {gpu}."
|
|
steps = "GPU уже задействован — ничего делать не нужно."
|
|
elif not built: # CPU-only torch build — the usual case
|
|
if gpus:
|
|
summary = (
|
|
"Главная причина: установлена CPU-сборка PyTorch (+cpu) — она физически "
|
|
f"не умеет в CUDA. Видеокарта ({gpus[0]}) и драйвер {driver or '?'} на месте, "
|
|
"поэтому достаточно переустановить PyTorch со сборкой CUDA."
|
|
)
|
|
steps = ("Переустановите PyTorch под CUDA, затем перезапустите приложение:\n\n"
|
|
+ command)
|
|
else:
|
|
summary = (
|
|
"Установлена CPU-сборка PyTorch (+cpu), и видеокарта NVIDIA не обнаружена "
|
|
"(nvidia-smi не отвечает). Либо нет NVIDIA GPU, либо не установлен драйвер."
|
|
)
|
|
steps = ("1. Проверьте видеокарту и драйвер: в консоли выполните nvidia-smi\n"
|
|
"2. Если NVIDIA GPU есть — переустановите PyTorch под CUDA:\n\n"
|
|
+ command +
|
|
"\n\nБез NVIDIA GPU всё работает на CPU — просто медленнее.")
|
|
else: # CUDA-enabled torch build, but CUDA still not available
|
|
if not gpus:
|
|
summary = (
|
|
f"PyTorch собран под CUDA {built}, но видеокарта/драйвер NVIDIA не найдены. "
|
|
"Скорее всего не установлен драйвер NVIDIA или нет GPU."
|
|
)
|
|
steps = "Установите свежий драйвер NVIDIA и перезапустите. Проверка: nvidia-smi."
|
|
elif cuda_driver and _ver_tuple(cuda_driver) < _ver_tuple(built):
|
|
summary = (
|
|
f"Драйвер поддерживает CUDA {cuda_driver}, а PyTorch собран под CUDA {built} — "
|
|
"версия драйвера слишком старая."
|
|
)
|
|
steps = ("Вариант A — обновите драйвер NVIDIA (рекомендуется).\n"
|
|
"Вариант B — поставьте PyTorch под CUDA вашего драйвера:\n\n" + command)
|
|
else:
|
|
summary = (
|
|
f"PyTorch собран под CUDA {built}, GPU ({gpus[0]}) есть, но CUDA недоступна — "
|
|
"вероятен конфликт версий или повреждённая установка."
|
|
)
|
|
steps = "Переустановите PyTorch под CUDA:\n\n" + command
|
|
return {"summary": summary, "details": details, "steps": steps, "command": command}
|