python实现nvidia-smi功能

#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2021/11/24
# @Author  : linzhe

import pynvml
pynvml.nvmlInit()
gpunum= pynvml.nvmlDeviceGetCount()
for i in range(gpunum):
    try:
        handle = (pynvml.nvmlDeviceGetHandleByIndex(i)) # GPU id
        meminfo = pynvml.nvmlDeviceGetMemoryInfo(handle)
        print('gpu%d \t used:%5d/%d \t free:%d/%d'%(i, meminfo.used//1048576, meminfo.total//1048576, meminfo.free//1048576, meminfo.total//1048576))
    except pynvml.NVMLError_GpuIsLost:
        print('gpu%d lost'%i)

效果

gpu0     used:    3/11019        free:11016/11019
gpu1     used:10888/11019        free:131/11019
gpu2     used:10888/11019        free:131/11019
gpu3     used:    3/11019        free:11016/11019
gpu4 lost
gpu5     used: 3680/11019        free:7339/11019
gpu6     used: 9580/11019        free:1439/11019