Pytorch手动释放显存

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


torch.cuda.empty_cache()

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

发表回复

您的邮箱地址不会被公开。 必填项已用 * 标注