转自:https://oldpan.me/archives/pytorch-gpu-memory-usage-track
Pytorch已经可以自动回收我们不用的显存,类似于python的引用机制,当某一内存内的数据不再有任何变量引用时,这部分的内存便会被释放。但有一点需要注意,当我们有一部分显存不再使用的时候,这部分释放的显存通过Nvidia-smi命令是看不到的,举个例子:
device
= torch
.device
('cuda:0')
dummy_tensor_4
= torch
.randn
(120, 3, 512, 512).float().to
(device
)
dummy_tensor_5
= torch
.randn
(80, 3, 512, 512).float().to
(device
)
dummy_tensor_4
= dummy_tensor_4
.cpu
()
dummy_tensor_2
= dummy_tensor_2
.cpu
()
torch
.cuda
.empty_cache
()
Pytorch的开发者也对此进行说明了,这部分释放后的显存可以用,只不过不在Nvidia-smi中显示罢了。