众所周知,Pytorch中简单地del某个tensor并不能释放掉该tensor占用的显存,还需要搭配torch.cuda.empty()来进行。但是,这也有例外。比如,如果该tensor在一个依然alive的计算图中,就无法被释放显存。下面是示例:
import torchdef active_bytes():stats = torch.cuda.memory_stats()current_active_byte = stats['active_bytes.all.current']return current_active_byte# initial usage
print("Init usage {}". format(active_bytes()))# vanilla tensor
x = torch.randn((256, 128), device='cuda')
w = torch.randn((128, 512), device='cuda')
l = torch.matmul(x, w)
print("Vanilla tensor {}". format(active_bytes()))del x
print("Vanilla tensor: del x {}". format(active_bytes()))
del w
print("Vanilla tensor: del w {}". format(active_bytes()))
l = l.cpu()
print("Vanilla tensor: l = l.cpu() {}". format(active_bytes()))# requires_grad=True
x = torch.randn((256, 128), device='cuda', requires_grad=True)
w = torch.randn((128, 512), device='cuda', requires_grad=True)
l = torch.matmul(x, w)
print("requires_grad=True {}". format(active_bytes()))del x
print("requires_grad=True: del x {}". format(active_bytes()))
del w
print("requires_grad=True: del w {}". format(active_bytes()))
l = l.cpu()
print("requires_grad=True: l = l.cpu() {}". format(active_bytes()))
以上代码的运行结果如下:
Init usage 0
Vanilla tensor 917504
Vanilla tensor: del x 786432
Vanilla tensor: del w 524288
Vanilla tensor: l = l.cpu() 0
requires_grad=True 917504
requires_grad=True: del x 917504
requires_grad=True: del w 917504
requires_grad=True: l = l.cpu() 393216
可以看到:在设置叶子节点x、w的requires_grad=True之后,del并不能释放x、w所占用的显存。我的理解是:w、x均需参与backward的过程,属于计算图的一部分,而此时计算图仍然“alive”,所以无法释放w、x的显存。那么,我把计算图整个移走可不可以呢?答案是:不行。Pytorch的计算图一旦初始化(即调用了forward),就不会改变位置。也就是说,如果你在某块GPU上进行了forward操作,你就只能在这块GPU上进行backward。
这种机制导致了一个很现实的状况:你在训练模型前最好提前估算好模型占用的显存规模,乖乖地选择一块显存足够用的GPU(nvidia狂喜)。