1.获取gpu信息的文件gpu_info.py
cat gpu_info.py
import pynvml #获取GPU信息
class GpuInfo(object):
def __init__(self):
#初始化
pynvml.nvmlInit()
def get_gpu_device(self):
deviceCount = pynvml.nvmlDeviceGetCount()
gpu_list = []
for i in range(deviceCount):
handle = pynvml.nvmlDeviceGetHandleByIndex(i)
print("GPU", i, ":", pynvml.nvmlDeviceGetName(handle))
gpu_list.append(i)
return gpu_list
def get_free_rate(self, gpu_id):
handle = pynvml.nvmlDeviceGetHandleByIndex(gpu_id)
info = pynvml.nvmlDeviceGetMemoryInfo(handle)
free_rate = int((info.free / info.total) * 100)
return free_rate
def get_gpu_info(self, gpu_id):
handle = pynvml.nvmlDeviceGetHandleByIndex(gpu_id)
info = pynvml.nvmlDeviceGetMemoryInfo(handle)
M = 1024*1024
gpu_info = "id:{} total:{}M free:{}M used:{}M free_rate:{}%".format(gpu_id, info.total/M, info.free/M, info.used/M, self.get_free_rate(gpu_id))
return gpu_info
def release(self):
#最后要关闭管理工具
pynvml.nvmlShutdown()
if __name__ == "__main__":
print("GPU INFO----------------------------")
gpu_info = GpuInfo()
gpu_devices = gpu_info.get_gpu_device()
print("GPU USE INFO----------------------------")
for gpuid in gpu_devices:
print(gpu_info.get_gpu_info(gpuid))
gpu_info.release()
2. 引用gpu_info.py文件
cat app.py
from gpu_info import GpuInfo
gf = GpuInfo()
a = gf.get_gpu_device()
print(a)