PyTorch：为什么运行 output = model(images) 使用这么多 GPU 内存？

Question

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

output = model(images)

其中 images 是一些输入张量， model 是我的 PyTorch 模型。在运行该行之前，我有大约 9GB 的可用 GPU 内存，之后我的可用内存降至 2.5GB（运行loss = criterion(outputs, labels).

两个问题：

1. 这是正常的吗？
2. 为什么会发生这种情况？这些内存都用来做什么？据我了解，模型已经加载，并且在进行调用之前实际的输入张量已经在 GPU 上。输出张量本身不可能那么大。它与存储计算图有关系吗？

Answer 1

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

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

您会发现使用的内存少得多，因为不会存储计算图。但这也意味着您无法再计算导数。然而，如果您只想评估模型而不需要任何优化，那么这是标准方法。

如果您仍然想优化，有一种方法可以减少内存消耗，它称为检查点。每当您在向后传递中需要中间张量时，都会从输入（或实际上从最后一个“检查点”）再次计算它，而不存储直到该张量的中间张量。但这只是计算成本更高。您正在用内存来交换计算时间。