【转载】 Pytorch手动释放显存

原文地址:

http://www.shijinglei.com/2020/04/20/pytorch%E9%87%8A%E6%94%BE%E6%98%BE%E5%AD%98/

===========================================

最近遇到一个问题,用pytorch跑一个 不固定输入的模型inference,一张图片一张图片的测试。有两张图片分辨率相同,都是4032×3024,但是前一张图片可以跑,到后一张图片就报cuda out of memory的错误。原因是对于固定输入,pytorch会复用显存。而对于不固定输入的模型,pytorch存在显存释放不及时的问题,导致一直增长。解决方法是在每次inference后,手动释放显存

torch.cuda.empty_cache()

加上该语句之后就可以正常运行了

=============================================