Преглед на файлове

修改为获取npu信息

zhy преди 2 седмици
родител
ревизия
b16ac75928
променени са 1 файла, в които са добавени 26 реда и са изтрити 3 реда
  1. 26 3
      tests/detection_performance_loss_test.py

+ 26 - 3
tests/detection_performance_loss_test.py

@@ -7,7 +7,7 @@ import argparse
 import os
 
 import psutil
-import GPUtil
+# import GPUtil
 import numpy as np
 import time
 from threading import Thread
@@ -19,6 +19,27 @@ from watermark_verify.inference.ssd_inference import SSDInference
 from watermark_verify.inference.yolox_inference import YOLOXInference
 from watermark_verify.tools.evaluate_tool import calculate_iou
 
+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:
@@ -39,9 +60,11 @@ 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[0].load * 100)  # 获取第一个 GPU 的使用率
+                self.gpu_usage.append(gpus)
             else:
                 self.gpu_usage.append(0)  # 若没有 GPU 则记为 0