我有一个由
tensorflow-gpu
运行的模型,我的设备是 nvidia
。我想列出每秒的 GPU 使用情况,以便可以测量平均/最大 GPU 使用情况。我可以通过打开两个终端来手动执行此操作,一个是运行模型,另一个是通过 nvidia-smi -l 1
进行测量。当然,这不是一个好办法。我也尝试使用 Thread
来做到这一点,就是这样。
import subprocess as sp
import os
from threading import Thread
class MyThread(Thread):
def __init__(self, func, args):
super(MyThread, self).__init__()
self.func = func
self.args = args
def run(self):
self.result = self.func(*self.args)
def get_result(self):
return self.result
def get_gpu_memory():
output_to_list = lambda x: x.decode('ascii').split('\n')[:-1]
ACCEPTABLE_AVAILABLE_MEMORY = 1024
COMMAND = "nvidia-smi -l 1 --query-gpu=memory.used --format=csv"
memory_use_info = output_to_list(sp.check_output(COMMAND.split()))[1:]
memory_use_values = [int(x.split()[0]) for i, x in enumerate(memory_use_info)]
return memory_use_values
def run():
pass
t1 = MyThread(run, args=())
t2 = MyThread(get_gpu_memory, args=())
t1.start()
t2.start()
t1.join()
t2.join()
res1 = t2.get_result()
但是,这也不会返回每秒的使用情况。有好的解决办法吗?
在命令中
nvidia-smi -l 1 --query-gpu=memory.used --format=csv
-l 代表:
-l, --loop= 探测,直到按指定的第二个间隔按下 Ctrl+C。
所以命令:
COMMAND = 'nvidia-smi -l 1 --query-gpu=memory.used --format=csv'
sp.check_output(COMMAND.split())
永远不会终止并返回。
如果将事件循环从命令(nvidia-smi)删除到python,它就会起作用。
这是代码:
import subprocess as sp
import os
from threading import Thread , Timer
import sched, time
def get_gpu_memory():
output_to_list = lambda x: x.decode('ascii').split('\n')[:-1]
ACCEPTABLE_AVAILABLE_MEMORY = 1024
COMMAND = "nvidia-smi --query-gpu=memory.used --format=csv"
try:
memory_use_info = output_to_list(sp.check_output(COMMAND.split(),stderr=sp.STDOUT))[1:]
except sp.CalledProcessError as e:
raise RuntimeError("command '{}' return with error (code {}): {}".format(e.cmd, e.returncode, e.output))
memory_use_values = [int(x.split()[0]) for i, x in enumerate(memory_use_info)]
# print(memory_use_values)
return memory_use_values
def print_gpu_memory_every_5secs():
"""
This function calls itself every 5 secs and print the gpu_memory.
"""
Timer(5.0, print_gpu_memory_every_5secs).start()
print(get_gpu_memory())
print_gpu_memory_every_5secs()
"""
Do stuff.
"""
尝试
pip install nvidia-ml-py3
:
import nvidia_smi
nvidia_smi.nvmlInit()
deviceCount = nvidia_smi.nvmlDeviceGetCount()
for i in range(deviceCount):
handle = nvidia_smi.nvmlDeviceGetHandleByIndex(i)
util = nvidia_smi.nvmlDeviceGetUtilizationRates(handle)
mem = nvidia_smi.nvmlDeviceGetMemoryInfo(handle)
print(f"|Device {i}| Mem Free: {mem.free/1024**2:5.2f}MB / {mem.total/1024**2:5.2f}MB | gpu-util: {util.gpu/100.0:3.1%} | gpu-mem: {util.memory/100.0:3.1%} |")
这是获得此输出的更基本的方法,但同样有效 - 而且我认为更容易理解。我添加了一个小型 10 值缓存以获得良好的近期平均值,并将检查时间提高到每秒。它输出最后 10 秒和当前每秒的平均值,因此可以识别导致使用的操作(我认为最初的问题是)。
import subprocess as sp
import time
memory_total=8192 #found with this command: nvidia-smi --query-gpu=memory.total --format=csv
memory_used_command = "nvidia-smi --query-gpu=memory.used --format=csv"
isolate_memory_value = lambda x: "".join(y for y in x.decode('ascii') if y in "0123456789")
def main():
percentage_cache = []
while True:
memory_used = isolate_memory_value(sp.check_output(memory_used_command.split(), stderr=sp.STDOUT))
percentage = float(memory_used)/float(memory_total)*100
percentage_cache.append(percentage)
percentage_cache = percentage_cache[max(0, len(percentage_cache) - 10):]
print("curr: " + str(percentage) + " %", "\navg: " + str(sum(percentage_cache)/len(percentage_cache))[:4] + " %\n")
time.sleep(1)
main()
您可能想使用
https://github.com/anderskm/gputil#usage中的
GPutil
# For prints
GPUtil.showUtilization()
# To get values
GPUs = GPUtil.getGPUs()
load = GPUs[0].load