PyTorch:为什么运行output =model(映像)需要这么多GPU内存？

Question

在试图理解为什么我的PyTorch模型的最大批处理大小是有限的时，我注意到使用最多内存的不是模型本身，也不是将张量加载到GPU中。当第一次生成预测时，大部分内存都会耗尽，例如训练循环中的以下一行：

output = model(images)

其中图像是一些输入张量，模型是我的PyTorch模型。在运行该行之前，我有大约9GB的GPU内存可用，然后我降低到2.5GB (在运行loss = criterion(outputs, labels)之后，它进一步下降到1GB可用)。

两个问题：

1. 这是正常的吗？
2. 为什么会这样？那些记忆是用来做什么的？据我所知，模型已经加载，在进行调用之前，实际的输入张量已经在GPU上。输出张量本身不可能那么大。它与存储计算图表有关吗？

Answer 1

这是正常的:这里的关键是所有中间张量(整个计算图)都必须被存储，如果你想通过后向模式微分计算梯度。您可以通过使用.no_grad上下文管理器来避免这种情况：

with torch.no_grad():
    output = model(images)

您将看到使用的内存要少得多，因为没有存储计算图表。但这也意味着你不能再计算导数了。但是，如果您只想评估模型而不需要任何优化，这是标准的方法。

如果您仍然想要优化，有一种方法可以减少内存的占用，这就是所谓的检查点。当你需要一个中间张量在向后通过，它将再次计算从输入(或实际上从最后的“检查点”)，而没有存储一个中间张量，直到那个张量。但这只是在计算上更昂贵。你用内存来交换计算时间。