File size: 2,762 Bytes
2979822 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 | import platform
import onnxruntime as ort
from typing import Optional
try:
import psutil
HAS_PSUTIL = True
except ImportError:
HAS_PSUTIL = False
try:
import cpuinfo
HAS_CPUINFO = True
except ImportError:
HAS_CPUINFO = False
try:
import GPUtil
HAS_GPUTIL = True
except ImportError:
HAS_GPUTIL = False
def get_cpu_info() -> str:
cpu_name = None
cpu_freq_mhz = None
cpu_cores = None
cpu_threads = None
if HAS_CPUINFO:
try:
info = cpuinfo.get_cpu_info()
cpu_name = (
info.get("brand_raw") or info.get("brand") or info.get("model name")
)
if cpu_name:
cpu_name = cpu_name.strip()
except Exception:
pass
if HAS_PSUTIL:
try:
cpu_freq = psutil.cpu_freq()
cpu_cores = psutil.cpu_count(logical=False)
cpu_threads = psutil.cpu_count(logical=True)
if cpu_freq and cpu_freq.current:
cpu_freq_mhz = cpu_freq.current
except Exception:
pass
if not cpu_name:
cpu_name = platform.processor() or "Unknown CPU"
parts = []
if cpu_name:
parts.append(cpu_name)
if cpu_cores and cpu_threads:
if cpu_cores == cpu_threads:
parts.append(f"{cpu_cores} cores")
else:
parts.append(f"{cpu_cores}C/{cpu_threads}T")
if cpu_freq_mhz:
if cpu_freq_mhz >= 1000:
parts.append(f"{cpu_freq_mhz/1000:.2f} GHz")
else:
parts.append(f"{cpu_freq_mhz:.0f} MHz")
return " | ".join(parts) if parts else "Unknown CPU"
def get_gpu_info() -> Optional[str]:
if HAS_GPUTIL:
try:
gpus = GPUtil.getGPUs()
if gpus and len(gpus) > 0:
gpu = gpus[0]
return gpu.name.strip()
except Exception:
pass
try:
import subprocess
result = subprocess.run(
["nvidia-smi", "--query-gpu=name", "--format=csv,noheader"],
capture_output=True,
text=True,
timeout=2,
)
if result.returncode == 0 and result.stdout.strip():
return result.stdout.strip().split("\n")[0]
except Exception:
pass
return None
def get_execution_provider_name(session: ort.InferenceSession) -> str:
try:
providers = session.get_providers()
if "CUDAExecutionProvider" in providers:
return "CUDA"
elif "CPUExecutionProvider" in providers:
return "CPU"
elif providers:
return providers[0].replace("ExecutionProvider", "")
return "Unknown"
except Exception:
return "Unknown"
|