Variable的backward()方法中的参数retain_graph是什么意思?

What does the parameter retain_graph mean in the Variable's backward() method?

我正在查看 neural transfer pytorch tutorial 并且对 retain_variable(已弃用,现在称为 retain_graph)的使用感到困惑。代码示例显示:

class ContentLoss(nn.Module):

    def __init__(self, target, weight):
        super(ContentLoss, self).__init__()
        self.target = target.detach() * weight
        self.weight = weight
        self.criterion = nn.MSELoss()

    def forward(self, input):
        self.loss = self.criterion(input * self.weight, self.target)
        self.output = input
        return self.output

    def backward(self, retain_variables=True):
        #Why is retain_variables True??
        self.loss.backward(retain_variables=retain_variables)
        return self.loss

来自the documentation

retain_graph (bool, optional) – If False, the graph used to compute the grad will be freed. Note that in nearly all cases setting this option to True is not needed and often can be worked around in a much more efficient way. Defaults to the value of create_graph.

因此,通过设置 retain_graph= True,我们不会释放在向后传递时为图形分配的内存。保留这段内存有什么好处,我们为什么需要它?

当您有多个网络输出时,这是一个非常有用的功能。这是一个完全虚构的例子:假设你想构建一些随机卷积网络,你可以问两个问题:输入图像是否包含猫,图像是否包含汽车?

这样做的一种方法是拥有一个共享卷积层的网络,但后面有两个平行的分类层(原谅我糟糕的 ASCII 图,但这应该是三个卷积层,后面是三个完全连接的层图层,一层用于猫,一层用于汽车):

                    -- FC - FC - FC - cat?
Conv - Conv - Conv -|
                    -- FC - FC - FC - car?

给定一张我们想要 运行 两个分支的图片,在训练网络时,我们可以通过多种方式进行。首先(这可能是这里最好的事情,说明这个例子有多糟糕),我们简单地计算两个评估的损失并对损失求和,然后反向传播。

但是,还有另一种情况 - 我们希望按顺序执行此操作。首先我们想通过一个分支进行反向传播,然后通过另一个分支进行反向传播(我之前有过这个用例,所以它不是完全编造的)。在那种情况下,一个图上的 运行ning .backward() 也会破坏卷积层中的任何梯度信息,并且第二个分支的卷积计算(因为这些是唯一与另一个分支共享的)将不再包含图表!这意味着,当我们尝试通过第二个分支进行反向传播时,Pytorch 将抛出错误,因为它找不到将输入连接到输出的图! 在这些情况下,我们可以通过在第一次向后传递时简单地保留图形来解决问题。然后该图将不会被消耗,而只会被不需要保留它的第一个反向传递所消耗。

编辑:如果您在所有向后传递中保留图形,则永远不会释放附加到输出变量的隐式图形定义。这里也可能有一个用例,但我想不出一个。所以一般来说,你应该确保最后一次向后传递通过不保留图形信息来释放内存。

至于多次向后传递会发生什么:正如您猜到的那样,pytorch 通过就地添加它们来累积梯度(到变量's/parameters .grad 属性)。 这可能非常有用,因为这意味着循环批处理并一次处理一次,在最后累积梯度,将执行与执行完整批处理更新相同的优化步骤(仅将所有梯度汇总为出色地)。虽然完全批处理的更新可以更多地并行化,因此通常更可取,但在某些情况下,批处理计算要么非常非常难以实现,要么根本不可能实现。然而,使用这种积累,我们仍然可以依赖批处理带来的一些很好的稳定特性。 (如果不是性能增益)

@cleros 在 retain_graph=True 的使用上讲得很到位。本质上,它会保留计算某个变量所需的任何信息,以便我们可以对其进行反向传递。

说明性示例

假设我们有一个如上所示的计算图。变量de是输出,a是输入。例如,

import torch
from torch.autograd import Variable
a = Variable(torch.rand(1, 4), requires_grad=True)
b = a**2
c = b*2
d = c.mean()
e = c.sum()

当我们 d.backward() 时,那很好。计算完成后,图中计算 d 的部分将默认释放以节省内存。所以如果我们做e.backward(),就会弹出错误信息。为了做到e.backward(),我们必须在d.backward()中将参数retain_graph设置为True,即

d.backward(retain_graph=True)

只要你在向后的方法中使用retain_graph=True,你可以随时向后做:

d.backward(retain_graph=True) # fine
e.backward(retain_graph=True) # fine
d.backward() # also fine
e.backward() # error will occur!

可以找到更多有用的讨论 here

真实用例

现在,一个真实的用例是多任务学习,其中您有多个损失,可能位于不同的层。假设您有 2 个损失:loss1loss2,它们位于不同的层。为了将 loss1loss2 w.r.t 的梯度独立地反向传播到网络的可学习权重。您必须在第一个反向传播损失中使用 backward() 方法中的 retain_graph=True

# suppose you first back-propagate loss1, then loss2 (you can also do the reverse)
loss1.backward(retain_graph=True)
loss2.backward() # now the graph is freed, and next process of batch gradient descent is ready

optimizer.step() # update the network parameters