Python 查看显存大小

版权声明：本文为博主原创文章，遵循 CC 4.0 BY-SA 版权协议，转载请附上原文出处链接和本声明。

本文链接： https://blog.csdn.net/u014421797/article/details/89875821

背景：因为一些需求，需要在python程序中查看显存的占用情况。

1.安装NVML库

pip install nvidia-ml-py

2.使用

import pynvml
pynvml.nvmlInit()
handle = pynvml.nvmlDeviceGetHandleByIndex(0)   #gpu_id
meminfo = pynvml.nvmlDeviceGetMemoryInfo(handle)
print(meminfo.used)

结果会显示显卡0的显存占用情况。