如果我只对一些样本进行转发，什么时候释放计算图？

Question

我有一个用例，我对批次中的每个样本进行转发，并且仅根据样本模型输出的某些条件累积某些样本的损失。这是一个说明代码，

for batch_idx, (data, target) in enumerate(train_loader):
    optimizer.zero_grad()
    total_loss = 0

    loss_count_local = 0
    for i in range(len(target)):
        im = Variable(data[i].unsqueeze(0).cuda())
        y = Variable(torch.FloatTensor([target[i]]).cuda())

        out = model(im)

        # if out satisfy some condtion, we will calculate loss
        # for this sample, else proceed to next sample
        if some_condition(out):
            loss = criterion(out, y)
        else:
            continue

        total_loss += loss
        loss_count_local += 1

        if loss_count_local == 32 or i == (len(target)-1):
            total_loss /= loss_count_local
            total_loss.backward()
            total_loss = 0
            loss_count_local = 0

    optimizer.step()

我的问题是，正如我对所有样本所做的那样，但只对某些样本进行后退。那些对损失没有贡献的样本的图形何时会被释放？这些图是仅在 for 循环结束后释放还是在我转发下一个样本后立即释放？我在这里有点困惑。

同样对于那些确实对 total_loss 有贡献的样本，它们的图表将在我们 total_loss.backward() 之后立即释放。是吗？

Answer 1

让我们从 PyTorch 如何释放内存的一般性讨论开始：

首先，我们应该强调 PyTorch 使用存储在 Python 对象属性中的隐式声明图。（记住，它是 Python，所以一切都是对象）。更具体地说，torch.autograd.Variable 具有 .grad_fn 属性。这个属性的类型定义了我们有什么样的计算节点（例如加法），以及该节点的输入。

这很重要，因为 Pytorch 只需使用标准 python 垃圾收集器（如果相当积极的话）即可释放内存。在此上下文中，这意味着（隐式声明的）计算图将保持活动状态，只要在当前范围内存在对保存它们的对象的引用！

这意味着如果您对样本 s_1 ... s_k 进行某种批处理，计算每个样本的损失并在末尾添加损失，累积损失将包含对每个单独损失的引用，而这些损失又包含引用到计算它的每个计算节点。

所以你的代码问题更多是关于 Python（或者更具体地说，它的垃圾收集器）如何处理引用，而不是 Pytorch。由于您在一个对象 (total_loss) 中累积了损失，因此您使指针保持活动状态，因此在您在外循环中重新初始化该对象之前不会释放内存。

应用于您的示例，这意味着您在正向传递（在 out = model(im) 处）创建的计算图仅由 out 对象及其任何未来计算引用。因此，如果您计算损失并对其求和，您将保持对 out 的引用，从而保持对计算图的引用。但是，如果您不使用它，垃圾收集器应该递归收集 out 及其计算图。

如果我只对一些样本进行转发，什么时候释放计算图？

When will the computation graph be freed if I only do forward for some samples?

python

backpropagation

pytorch