如何在 Python 中获取每一秒的 GPU 使用率

Question

我有一个型号运行由 tensorflow-gpu 生产，我的设备是 nvidia。我想列出每一秒的 GPU 使用情况，以便我可以测量 average/max 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()

但是，这并不 return 每一秒的使用量。有什么好的解决办法吗？

Answer 1

在命令中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())

永远不会终止并且 return。

如果您从命令 (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.
"""

如何在 Python 中获取每一秒的 GPU 使用率

How to get every second's GPU usage in Python

python

nvidia