Python 获取系统 GPU 占用率及指定进程 GPU 占用率并写入 CSV 文件

要获取系统GPU占用率，可以使用第三方库pynvml。首先需要安装该库：

pip install pynvml

然后可以使用以下代码获取系统GPU占用率并写入csv文件：

import csv
import time
import pynvml

# 初始化pynvml
pynvml.nvmlInit()

# 获取GPU数量
device_count = pynvml.nvmlDeviceGetCount()

# 获取所有GPU的占用率
def get_gpu_usage():
    gpu_usage = []
    for i in range(device_count):
        handle = pynvml.nvmlDeviceGetHandleByIndex(i)
        info = pynvml.nvmlDeviceGetUtilizationRates(handle)
        usage = info.gpu
        gpu_usage.append(usage)
    return gpu_usage

# 获取指定进程的GPU占用率
def get_process_gpu_usage(process_name):
    gpu_usage = []
    for i in range(device_count):
        handle = pynvml.nvmlDeviceGetHandleByIndex(i)
        # 获取指定进程的显存占用
        info = pynvml.nvmlDeviceGetMemoryInfo(handle)
        # 获取所有正在运行的进程
        processes = pynvml.nvmlDeviceGetComputeRunningProcesses(handle)
        for p in processes:
            # 根据进程名字获取指定进程的GPU占用率
            if pynvml.nvmlProcessGetName(p.pid).decode() == process_name:
                gpu_usage.append(p.usedGpuMemory / info.total)
                break
        else:
            gpu_usage.append(0)
    return gpu_usage

# 写入.csv文件
def write_to_csv(data, filename):
    with open(filename, 'w', newline='') as csvfile:
        writer = csv.writer(csvfile)
        writer.writerow(["GPU Index", "GPU Usage"])
        for i, usage in enumerate(data):
            writer.writerow([i, usage])

# 获取系统GPU占用率
system_gpu_usage = get_gpu_usage()
write_to_csv(system_gpu_usage, "system_gpu_usage.csv")

# 获取指定进程的GPU占用率
process_name = "your_process_name"
process_gpu_usage = get_process_gpu_usage(process_name)
write_to_csv(process_gpu_usage, "process_gpu_usage.csv")

# 关闭pynvml
pynvml.nvmlShutdown()

请将your_process_name替换为你要监视的进程名字。运行上述代码后，将会生成system_gpu_usage.csv和process_gpu_usage.csv两个文件，分别包含系统GPU占用率和指定进程的GPU占用率。