小编Viv*_*tel的帖子

我有一个运行的模型，tensorflow-gpu我的设备是nvidia. 我想列出每秒的 GPU 使用情况，以便可以测量平均/最大 GPU 使用情况。我可以通过打开两个终端来手动执行此操作，一个是运行模型，另一个是通过nvidia-smi -l 1. 当然，这不是一个好办法。我也尝试使用 aThread来做到这一点，就是这样。

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 …