Просмотр исходного кода

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

zhy 2 недель назад
Родитель
Сommit
74e4550f0e
1 измененных файлов с 26 добавлено и 3 удалено
  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 os
 
 
 import psutil
 import psutil
-import GPUtil
+# import GPUtil
 import numpy as np
 import numpy as np
 import time
 import time
 from threading import Thread
 from threading import Thread
 import onnxruntime as ort
 import onnxruntime as ort
 from PIL import Image
 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:
 class UsageMonitor:
@@ -34,9 +56,10 @@ class UsageMonitor:
             self.cpu_usage.append(psutil.cpu_percent(interval=None))
             self.cpu_usage.append(psutil.cpu_percent(interval=None))
 
 
             # 记录 GPU 使用率
             # 记录 GPU 使用率
-            gpus = GPUtil.getGPUs()
+            # gpus = GPUtil.getGPUs()
+            gpus = get_npu_usage()
             if gpus:
             if gpus:
-                self.gpu_usage.append(gpus[0].load * 100)  # 获取第一个 GPU 的使用率
+                self.gpu_usage.append(gpus)  # 获取第一个 GPU 的使用率
             else:
             else:
                 self.gpu_usage.append(0)  # 若没有 GPU 则记为 0
                 self.gpu_usage.append(0)  # 若没有 GPU 则记为 0