当前位置:   article > 正文

pytorch做inference的时候如何减少显存占用_inference cache

inference cache

1、采用with torch.no_grad():

  1. with torch.no_grad():
  2. tc_data = netG_A(input_data, d_f)

相比于直接做inference,加上with torch.no_grad()能够占用更少的显存。即使设置了netG_A.eval(),占用显存和没设置时是一样的。难道eval()模式也还会保留中间梯度?

tc_data = netG_A(input_data, d_f)

2、删除中间变量:

  1. with torch.no_grad():
  2. tc_data = netG_A(input_data, d_f)
  3. c_data = tc_data.detach().clone()
  4. del tc_data
  5. torch.cuda.empty_cache()

由于我们只需要最终的结果,而模型inference的时候应该会保留一些中间变量,所以把预测结果的数据单独拷贝出来,然后再删掉原有的结果,能够减少显存占用。

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

闽ICP备14008679号