如何在Pytorch中查看内存中计算图的大小?

Hor*_*rse 5 python pytorch autograd

如何跟踪 cpu 上的前向传播创建的 autograd 图的内存分配?例如尝试在 cpu 上使用tracemalloc:

rnn=nn.RNNCell(100,100).to('cuda')
x=torch.ones((1000,100),device='cuda')
tracemalloc.start(25)
while True:
   print(tracemalloc.get_traced_memory())
   x=rnn(x)
Run Code Online (Sandbox Code Playgroud)

随着图形在每个循环步骤中的增加,打印的内存应该不断增加,但是打印的内存来自

tracemalloc.get_traced_memory()
Run Code Online (Sandbox Code Playgroud)

第三次循环后保持不变。到底是怎么回事?