当前位置:   article > 正文

pytorch的显存机制torch.cuda.empty_cache()_orch.cuda.empty_cache() cpu

orch.cuda.empty_cache() cpu

参考  pytorch的显存机制torch.cuda.empty_cache() - 云+社区 - 腾讯云

Pytorch已经可以自动回收我们不用的显存,类似于python的引用机制,当某一内存内的数据不再有任何变量引用时,这部分的内存便会被释放。但有一点需要注意,当我们有一部分显存不再使用的时候,这部分释放的显存通过Nvidia-smi命令是看不到的,举个例子:

  1. device = torch.device('cuda:0')
  2. # 定义两个tensor
  3. dummy_tensor_4 = torch.randn(120, 3, 512, 512).float().to(device) # 120*3*512*512*4/1000/1000 = 377.48M
  4. dummy_tensor_5 = torch.randn(80, 3, 512, 512).float().to(device) # 80*3*512*512*4/1000/1000 = 251.64M
  5. # 然后释放
  6. dummy_tensor_4 = dummy_tensor_4.cpu()
  7. dummy_tensor_2 = dummy_tensor_2.cpu()
  8. # 这里虽然将上面的显存释放了,但是我们通过Nvidia-smi命令看到显存依然在占用
  9. torch.cuda.empty_cache()
  10. # 只有执行完上面这句,显存才会在Nvidia-smi中释放

Pytorch的开发者也对此进行说明了,这部分释放后的显存可以用,只不过不在Nvidia-smi中显示罢了。
在这里插入图片描述

声明:本文内容由网友自发贡献,不代表【wpsshop博客】立场,版权归原作者所有,本站不承担相应法律责任。如您发现有侵权的内容,请联系我们。转载请注明出处:https://www.wpsshop.cn/w/2023面试高手/article/detail/127399
推荐阅读
相关标签
  

闽ICP备14008679号