python - torch.einsum 的内存使用

Question

我一直在尝试调试某个模型，该模型在重复几次的层中使用 torch.einsum 运算符。

在尝试分析训练期间模型的 GPU 内存使用情况时，我注意到某个 Einsum 操作显着增加了内存使用量。我正在处理多维矩阵。操作是 torch.einsum('b q f n, b f n d -> b q f d', A, B)。

另外值得一提的是:

• x 之前被分配给一个相同形状的张量。
• 在此操作后的每一层(它们都是相同的)GPU内存线性增加，并且直到模型迭代结束才释放。

我一直想知道为什么这个操作会使用这么多内存，以及为什么在该层类型的每次迭代后内存都保持分配状态。

Answer 1

变量“x”确实被覆盖了，但张量数据保存在内存中(也称为层的激活)以供以后在反向传递中使用。

因此，您反过来有效地为 torch.einsum 的结果分配了新的内存数据，但您不会替换 x 的内存，即使它有似乎被覆盖了。

---

要将其传递给测试，您可以计算 torch.no_grad() 下的正向传递上下文管理器(这些激活不会保存在内存中)并查看与标准推理相比的内存使用差异。