Pytorch清空显存缓冲区(torch.cuda.empty_cache)

hxxjxw

已于 2022-09-26 16:18:55 修改

阅读量4.5w

点赞数 33

文章标签： pytorch

于 2021-08-18 12:09:47 首次发布

本文链接：https://blog.csdn.net/hxxjxw/article/details/119777443

版权

本文探讨了PyTorch中GPU显存的管理，特别是`torch.cuda.empty_cache()`函数的作用。该函数用于清空CUDA缓存，防止已释放的显存被旧数据占用。当程序运行过程中遇到显存不足问题时，适时调用此函数可以避免显存爆满。虽然这可能导致短暂的性能下降，但在某些情况下（如循环中）能有效防止显存溢出。建议在每个epoch开始时使用，以释放不必要的显存资源。

摘要由CSDN通过智能技术生成

torch.cuda.empty_cache()
因为PyTorch是有缓存区的设置的，意思就是一个Tensor就算被释放了，进程也不会把空闲出来的显存还给GPU，而是等待下一个Tensor来填入这一片被释放的空间。所以我们用nvidia-smi/gpustat看到的显存占用不会减少

用torch.cuda.empty_cache可以清空缓冲区

在程序中加上这句会使速度变慢一些，但是有些情况下会有用，例如程序之前test的时候总是爆显存，然后在循环中加上了这句就不爆了
    for i, data in enumerate(data_loader):
        torch.cuda.empty_cache()
        img_meta = data['img_meta'][0].data[0]
        img_name = img_meta[0]['filename'].split('/')[-1]
        with torch.no_grad():
            result = model(return_loss=False, rescale=not show, **data)
如果显存资源比较紧缺，可以在每个epoch开始时释放下不用的显存资源。
 torch.cuda.empty_cache()  # 释放显存