torch.cuda.empty_cache 清空显存

empty_cache()不会增加Pytorch可用的GPU内存量。但是,在某些情况下,它可能有助于减少GPU记忆的碎片化。有关GPU内存管理的更多详细信息,请参见内存管理。

因为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() # 释放显存
————————————————

发表评论

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