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及其计算图。

如果我只针对某些样本进行转发，何时会释放计算图？

1 个答案: