Pytorch清空显存缓冲区(torch.cuda.empty_cache)
torch.cuda.empty_cache()写代码的好习惯是在程序最后执行以下这句话
·
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() # 释放显存
欢迎来到FlagOS开发社区,这里是一个汇聚了AI开发者、数据科学家、机器学习爱好者以及业界专家的活力平台。我们致力于成为业内领先的Triton技术交流与应用分享的殿堂,为推动人工智能技术的普及与深化应用贡献力量。
更多推荐

所有评论(0)