本地知识库（2）：我在闲鱼买了台服务器，进行ai训练UE5

#!/usr/bin/env python3
import torch
import platform
import subprocess
import re

def get_driver_version():
    try:
        result = subprocess.check_output(["nvidia-smi"], encoding="utf-8")
        match = re.search(r"Driver Version:\s*([\d.]+)", result)
        return match.group(1) if match else "Unknown"
    except Exception:
        return "NVIDIA driver not found"

def get_cuda_version_from_driver():
    try:
        result = subprocess.check_output(["nvidia-smi"], encoding="utf-8")
        match = re.search(r"CUDA Version:\s*([\d.]+)", result)
        return match.group(1) if match else "Unknown"
    except Exception:
        return "Unknown"

def main():
    print("="*60)
    print("🔍 PyTorch + GPU 环境检测")
    print("="*60)
    print(f"🧠 Python 版本: {platform.python_version()}")
    print(f"💻 系统: {platform.system()} {platform.release()}")
    print(f"🐍 PyTorch 版本: {torch.__version__}")

    # 检查 CUDA 是否可用
    cuda_available = torch.cuda.is_available()
    print(f"⚙️ CUDA 是否可用: {cuda_available}")

    if cuda_available:
        print(f"📦 PyTorch 内置 CUDA 版本: {torch.version.cuda}")
        print(f"💽 GPU 数量: {torch.cuda.device_count()}")
        for i in range(torch.cuda.device_count()):
            print(f"   ├─ GPU[{i}]: {torch.cuda.get_device_name(i)}")
            cap = torch.cuda.get_device_capability(i)
            print(f"   │   Compute Capability: {cap[0]}.{cap[1]}")
        print(f"🧩 当前 GPU 设备: {torch.cuda.current_device()}")
    else:
        print("⚠️ 未检测到可用的 CUDA 设备。")

    # 驱动信息
    print("="*60)
    print("🧰 NVIDIA 驱动信息")
    driver_ver = get_driver_version()
    cuda_driver_ver = get_cuda_version_from_driver()
    print(f"   Driver Version: {driver_ver}")
    print(f"   CUDA Runtime (from driver): {cuda_driver_ver}")

    # 推荐版本判断
    try:
        cap = torch.cuda.get_device_capability(0)
        major, minor = cap
        print("="*60)
        print("🔎 兼容性分析")
        if major == 6:
            print("✅ 检测到 Pascal 架构 (如 Tesla P100)。")
            print("   推荐使用 PyTorch 2.0 ~ 2.5 + CUDA 11.8~12.4")
        elif major == 7:
            print("✅ 检测到 Volta / Turing 架构。")
            print("   推荐使用 PyTorch 2.1+ + CUDA 12.x")
        elif major >= 8:
            print("✅ 检测到 Ampere / Ada 架构。")
            print("   可使用最新 PyTorch + CUDA 12.x")
        else:
            print("⚠️ 未识别架构，可能过旧或驱动异常。")
    except Exception:
        print("无法分析 GPU 架构（可能无 CUDA 设备）。")

    print("="*60)
    print("✅ 检测完成")

if __name__ == "__main__":
    main()