获取GPU资源

import time
import csv
import pynvml
import argparse


def gpu_usage(gpu_info_csv, query_interval, max_time):
    header = ['time']
    pynvml.nvmlInit()
    gpu_num = pynvml.nvmlDeviceGetCount()  # 显示有几块GPU

    for i in range(gpu_num):
        header.append("gpu_{}".format(i))

    # write header
    with open(gpu_info_csv, 'w', encoding='utf-8', newline='') as fp:
        writer = csv.DictWriter(fp, fieldnames=header)
        writer.writeheader()

        max_query_time = int(max_time / query_interval)
        count = 0
        time_records = []
        message = ""
        while count < max_query_time:
            current_time = time.ctime()
            record = [current_time]
            for i in range(gpu_num):
                handle = pynvml.nvmlDeviceGetHandleByIndex(i)
                meminfo = pynvml.nvmlDeviceGetMemoryInfo(handle)
                gpu_usage = meminfo.used / 1024.0 / 1024.0
                record.append(gpu_usage)
                message += "gpu_{} usage: {}M, ".format(i, gpu_usage) 

            time_records.append(record)
            message = "time: {}".format(current_time) + message  
            print(message)
            
            writer.writerow(dict(zip(header, record)))
            
            message = " "
            time.sleep(query_interval)
            count += 1

if __name__ == '__main__':

    long_description = "Get gpu memory usage"
    parser = argparse.ArgumentParser(description=long_description)
    parser.add_argument('-o', '--output', nargs='?', default='/home/zxj/gpu_info.csv',
                        help='输出的csv路径')
    parser.add_argument('--interval', type=float, default=0.5, help='查询显存使用量的时间间隔')
    parser.add_argument('--max_time', type=int, default=5, help='程序运行最大时间（秒）')

    args = parser.parse_args()
    gpu_usage(args.output, args.interval, args.max_time)