Browse Source

修改为获取npu显存占用计算gpu使用率

zhy 2 weeks ago
parent
commit
74e4550f0e
1 changed files with 26 additions and 3 deletions
  1. 26 3
      tests/classification_performance_loss_test.py

+ 26 - 3
tests/classification_performance_loss_test.py

@@ -7,13 +7,35 @@ import argparse
 import os
 
 import psutil
-import GPUtil
+# import GPUtil
 import numpy as np
 import time
 from threading import Thread
 import onnxruntime as ort
 from PIL import Image
 
+import subprocess
+import re
+
+def get_npu_usage():
+    try:
+        output = subprocess.check_output("npu-smi info", shell=True, encoding="utf-8")
+        percent = None
+        for line in output.splitlines():
+            line = line.strip()
+            # 提取 like "1524 / 44280" 值大于0的行
+            m = re.search(r"([1-9]\d*)\s*/\s*([1-9]\d*)", line)
+            if m:
+                used = int(m.group(1))
+                total = int(m.group(2))
+                percent = used * 100.0 / total if total > 0 else 0
+                return percent
+        print("未找到NPU 0 显存占用信息")
+        return None
+    except Exception as e:
+        print("获取NPU信息失败:", e)
+        return None
+
 
 # 定义监控函数
 class UsageMonitor:
@@ -34,9 +56,10 @@ class UsageMonitor:
             self.cpu_usage.append(psutil.cpu_percent(interval=None))
 
             # 记录 GPU 使用率
-            gpus = GPUtil.getGPUs()
+            # gpus = GPUtil.getGPUs()
+            gpus = get_npu_usage()
             if gpus:
-                self.gpu_usage.append(gpus[0].load * 100)  # 获取第一个 GPU 的使用率
+                self.gpu_usage.append(gpus)  # 获取第一个 GPU 的使用率
             else:
                 self.gpu_usage.append(0)  # 若没有 GPU 则记为 0