监控显卡显存（python代码）

一、前言

我和我同学的代码，分别占用14G显存，而显卡的显存只有24G，没有办法同时跑，

所以，他先跑，我的代码时刻监控显存的使用情况，只要显卡显存大于14G，我的代码就自动启动。

二、代码

import pynvml
import time




def get_gpu_mem_info(gpu_id=0):
    """
    根据显卡 id 获取显存使用信息, 单位 MB
    :param gpu_id: 显卡 ID
    :return: total 所有的显存，used 当前使用的显存, free 可使用的显存
    """

    pynvml.nvmlInit()
    if gpu_id < 0 or gpu_id >= pynvml.nvmlDeviceGetCount():
        print(r'gpu_id {} 对应的显卡不存在!'.format(gpu_id))
        return 0, 0, 0

    handler = pynvml.nvmlDeviceGetHandleByIndex(gpu_id)
    meminfo = pynvml.nvmlDeviceGetMemoryInfo(handler)
    total = round(meminfo.total / 1024 / 1024, 2)
    used = round(meminfo.used / 1024 / 1024, 2)
    free = round(meminfo.free / 1024 / 1024, 2)
    return total, used, free


if __name__ == '__main__':
    while True:
        gpu_mem_total, gpu_mem_used, gpu_mem_free = get_gpu_mem_info(gpu_id=0)
        # print(r'当前显卡显存使用情况：总共 {} MB， 已经使用 {} MB， 剩余 {} MB'
        #        .format(gpu_mem_total, gpu_mem_used, gpu_mem_free))
        if gpu_mem_free >= 14000:
            main()
            break
        time.sleep(10)

希望以上内容对你们有所帮助。