pytorch 训练过程中查看GPU显存占用情况

应用 pynvml 库

import pynvml
pynvml.nvmlInit()
handle = pynvml.nvmlDeviceGetHandleByIndex(0) # 0表示显卡标号

# 在每一个要查看的地方都要重新定义一个meminfo 
meminfo = pynvml.nvmlDeviceGetMemoryInfo(handle)
print(meminfo.total/1024**2) #总的显存大小
print(meminfo.used/1024**2)  #已用显存大小
print(meminfo.free/1024**2)  #剩余显存大小
# 单位是MB，如果想看G就再除以一个1024

另外一篇总结帖：

https://blog.csdn.net/qq_45066628/article/details/122840847
查看Python 代码的内存和模型显存消耗以及优化

还有一种简单的方式

但好像是看设备上tensors当前所占用的 GPU 显存（而不是整体的显存占用）

print("1:{}".format(torch.cuda.memory_allocated(0)))

https://blog.csdn.net/weixin_48759194/article/details/122336402