【Python】GPU内存监控脚本

#!/usr/bin/python
# -*- coding: utf-8 -*-

import time
import pynvml
from matplotlib import pyplot as plt
import matplotlib.backends.backend_tkagg


class GPUMonitor(object):
    def __init__(self, sleep_time):
        pynvml.nvmlInit()
        pynvml.nvmlSystemGetDriverVersion()
        self.GPUCounts = pynvml.nvmlDeviceGetCount()
        self.GPU_counts_list = [[]] * self.GPUCounts
        self.time = [[]] * self.GPUCounts
        self.sleep_time = sleep_time  # 秒

    def monitor(self):
        try:
            n = 0
            while True:
                GPUCount = 0
                # 读取GPU句柄
                handle = pynvml.nvmlDeviceGetHandleByIndex(GPUCount)
                # 读取GPU内存信息
                info = pynvml.nvmlDeviceGetMemoryInfo(handle)
                total = f'{(info.total / 1024 / 1024 / 1024):.2f}'
                used = f'{(info.used / 1024 / 1024 / 1024):.2f}'
                free = f'{(info.free / 1024 / 1024 / 1024):.2f}'
                print(self.logfile(">>>>>>正在监控第{}块GPU内存<<<<<<<\n"
                                   "脚本已运行{}秒\n专用GPU内存：{}G\n"
                                   "已使用专用CPU内存：{}G\n剩余专用GPU内存：{}G\n"
                                   .format(GPUCount, n, float(total), float(used), float(free))))
                self.GPU_counts_list[GPUCount].append(float(used))
                self.time[GPUCount].append(n)
                self.paint(self.time[0], self.GPU_counts_list[0])
                time.sleep(self.sleep_time)
                n += self.sleep_time

        except:
            plt.savefig("CPU内存使用量.png")
            pynvml.nvmlShutdown()

    def paint(self, x_list, y_list):
        plt.clf()
        plt.plot(x_list, y_list)
        plt.title("GPU Usage Monitoring")
        plt.ylabel("GPU dedicated memory /G")
        plt.xlabel("time/s")
        plt.pause(0.1)  # 暂停一秒
        plt.ioff()

    def logfile(self, text):
        with open('image.log', 'a+', encoding='utf-8') as f:
            t = time.strftime('%y-%m-%d %H:%M:%S')
            text = t + " " + text + '\n'
            f.write(text)
        f.close()
        return text

    def abnormal(self):
        length = len(self.GPU_counts_list[0])
        average = sum(self.GPU_counts_list)/length
        self.logfile("平均专用GPU占用为：{}G".format(average))
        plt.savefig("CPU内存使用量.png")


if __name__ == "__main__":
    while True:
        times = input("请输入监控间隔时间(整秒>0)，按回车键开启监控：")
        if times.isdigit():
            if int(times) > 0:
                break
    a = GPUMonitor(int(times))
    try:
        a.monitor()
    except:
        plt.savefig("CPU内存使用量.png")