删除中间张量会影响 PyTorch 中的计算图吗？

Question

为了释放内存，我想知道是否可以在我的模型的前向方法中删除中间张量。这是一个最小化的示例场景：

def forward(self, input):
    x1, x2 = input
    x1 = some_layers(x1)
    x2 = some_layers(x2)
    x_conc = torch.cat((x1,x2),dim=1)
    x_conc = some_layers(x_conc)
    return x_conc

基本上，模型通过两个单独的块传递两个张量，然后将结果连接起来。对连接的张量应用进一步的操作。如果我在创建 x_conc 后运行 del x1 和 del x2 会影响计算图吗？

Answer 1

如果您想稍后执行自动微分，PyTorch 会将 x1, x2 张量存储在计算图中。另外，请注意，使用 del 运算符删除张量是可行的，但您不会看到 GPU 内存减少。为什么？因为内存被释放但没有返回到设备。这是一种优化技术，从用户的角度来看，内存已被“释放”。也就是说，内存现在可用于生成新的张量。

因此，不建议删除张量来释放 GPU 内存。