查看显存占用

一.

使用pynvml包（实际显存占用，包含cache）

import pynvml

pynvml.nvmlInit()
handle = pynvml.nvmlDeviceGetHandleByIndex(0) # 0表示显卡标号
meminfo = pynvml.nvmlDeviceGetMemoryInfo(handle)
print(meminfo.total/1024**2) #总的显存大小
print(meminfo.used/1024**2)  #已用显存大小
print(meminfo.free/1024**2)  #剩余显存大小

二.

torch.cuda.memory_allocated（理论上显存占用，不包含cache）

print(torch.cuda.memory_allocated(device=0) / (1024 * 1024))#0表示显卡号

文档：https://pytorch.org/docs/master/generated/torch.cuda.memory_allocated.html?highlight=torch%20cuda%20memory_allocated#torch.cuda.memory_allocated