关于Pytorch中tensor的自动梯度求导与反向传播的理解

关于Pytorch中tensor的自动梯度求导与反向传播的理解

参考资料：

• 详解Pytorch 自动微分里的（vector-Jacobian product）；
• PyTorch 的 backward 为什么有一个 grad_variables 参数？；
• 【one way的pytorch学习笔记】(四)autograd的流程机制原理；
• pytorch中backward()函数详解；
• PyTorch 的 Autograd；

为了更好的计算梯度，torch提供了一个名为的tensor的数据结构，专门用于深度学习中计算梯度。但并不是每一个tensor都被赋予自动计算梯度的功能，需要同时满足以下两个条件：

• is_leaf：True；
• requires_grad：True.

这里有一个问题：既然已经有requires_grad属性可以解释一个tensor是否需要自动计算梯度，那为什么还需要is_leaf属性呢？

原因可能是为了节省内存占用。 比如我们创建了一个MLP网络，为了优化模型，我们需要更新网络的参数，具体就是模型的weight和bias。也就是说，在反向传播时，我们只需要保存weight和bias的梯度用于更新，至于保存其它的中间变量的梯度则完全没有意义。 在这种情况下，只需要设置weight和bias为叶子节点（is_leaf：True），其它变量为非叶子节点（is_leaf：False），就可以实现这个效果。在实际检查网络的weight和bias的is_leaf后，发现确实是True。

总的来说，is_leaf属性可以决定在反向传播时是否需要保存计算得到的梯度值。

• 在默认情况下，用户创建的tensor变量具有is_leaf；

# is_leaf: True, requires_grad: False
a = torch.tensor([1.0])
# is_leaf: True, requires_grad: True
b = torch.tensor([2.0], requires_grad=True)
# 以下语句会报错，因为torch只会给float类型的tensor计算梯度
b = torch.tensor([2], requires_grad=True)  # int

• 但是非用户创建的变量只会继承requires_grad，不会继承is_leaf；

x = torch.tensor([1.0], requires_grad=True)
y = torch.tensor([2.0])
# is_leaf: False, requires_grad: True
# grad_fn: MulBackward (计算梯度的后向函数)
z = x * y
z.backward()  # 计算并储存梯度值

根据上述代码，已知$z=xy$，可以分别计算这两个变量的梯度：
$$\frac{\partial z}{\partial x}=y=2,\frac{\partial z}{\partial y}=x=1$$
因为y的requires_grad为False，故只有x的grad得到了更新。

print(x.grad)  # tensor([2.])
print(y.grad)  # None

这里需要注意一点，torch反向传播过一次梯度后，计算图就会被释放。如果再次执行backward()，程序就会报错；如果需要多次反向传播梯度，可在第一次使用backward()时加上retain_graph=True。

x = torch.tensor([1.0], requires_grad=True)
y = torch.tensor([2.0])
z = x * y
z.backward()  
z.backward()  # 报错: RuntimeError: Trying to backward through the graph a second time, but the saved intermediate results have already been freed.

x = torch.tensor([1.0], requires_grad=True)
y = torch.tensor([2.0])
z = x * y
# retain_graph=True只能保持一次，如果要一直保持，则每次都需要添加
z.backward(retain_graph=True)  
print(x.grad)  # tensor([2.])
z.backward(retain_graph=True)  # 多次反向传播后的梯度会累积
print(x.grad)  # tensor([4.])

上文只是介绍了scalar对tensor的求导过程（z对于[x, y]的求导）。那么，如果是tensor对tensor的求导呢？事实上，pytorch并不允许这种情况发生，大概原因就是这会导致计算会很复杂。具体分析可以参考文首的第二篇文章。

x = torch.tensor([2., 3., 1.], requires_grad=True)
# y1 = x1 ** 2; y2 = x2 ** 2; y3 = x3 ** 2
y = x ** 2
y.backward()
# 报错(不允许tensor对tensor的求导)
# RuntimeError: grad can be implicitly created only for scalar outputs

那么该如何处理这种tensor对tensor求导的情况呢？

当涉及到多变量对多变量的导数时，我们需要借助雅可比矩阵的力量。以上述程序的情况为例，雅可比矩阵应为：
$$J=\left(\begin{array}{ccc}\frac{\partial y_1}{\partial x_1}& \frac{\partial y_1}{\partial x_2}& \frac{\partial y_1}{\partial x_3}\\ \frac{\partial y_2}{\partial x_1}& \frac{\partial y_2}{\partial x_2}& \frac{\partial y_2}{\partial x_3}\\ \frac{\partial y_3}{\partial x_1}& \frac{\partial y_3}{\partial x_2}& \frac{\partial y_3}{\partial x_3}\end{array}\right)=\left(\begin{array}{ccc}2x_1& 0& 0\\ 0& 2x_2& 0\\ 0& 0& 2x_3\end{array}\right)$$
求出对应的雅可比矩阵后，然后会把grad_tensors与其做点积运算，最终得到梯度值。通常会把grad_tensors设置为维度与x相同的元素均为1的向量。
$$g=v\cdot J=\left(\begin{array}{ccc}1& 1& 1\end{array}\right)\left(\begin{array}{ccc}2x_1& 0& 0\\ 0& 2x_2& 0\\ 0& 0& 2x_3\end{array}\right)=\left(\begin{array}{ccc}2x_1& 2x_2& 2x_3\end{array}\right)=\left(\begin{array}{ccc}4& 6& 2\end{array}\right)$$

x = torch.tensor([2., 3., 1.], requires_grad=True)
y = x ** 2
y.backward(torch.ones_like(x))
print(x.grad)  # [4, 6, 2]

可以把grad_tensors视为各个部分（$y_1,y_2,y_3$对$x_i$的偏导数）的权重，通常情况会把所有部分的权重都设置为1，但也可以设置为其它值。
$$g=v\cdot J=\left(\begin{array}{ccc}1& 3& 2\end{array}\right)\left(\begin{array}{ccc}2x_1& 0& 0\\ 0& 2x_2& 0\\ 0& 0& 2x_3\end{array}\right)=\left(\begin{array}{ccc}2x_1& 6x_2& 4x_3\end{array}\right)=\left(\begin{array}{ccc}4& 18& 4\end{array}\right)$$

x = torch.tensor([2., 3., 1.], requires_grad=True)
y = x ** 2
y.backward(torch.tensor([1, 3, 2]))
print(x.grad)  # [4, 18, 4]

补充：关于梯度的累计问题

在有些时候，我们需要累计多次的梯度数据后，再对网络参数进行更新。那么，在pytorch中具体该如何实现的梯度的累计呢？

其实，在上文中已经提到过，计算图进行反向传播一次后，那些叶子节点的梯度信息就会被储存在grad中。在进行第二次反向传播时，之前被储存的梯度信息是不会被消失的。因此，在训练模型时，常用的一种做法就是在进行反向传播之前，使用zero_grad()清除之前的梯度。

但是，这里我们需要计算多次的反向传播得到的梯度之和，因此不必使用zero_grad()，而是直接进行反向传播即可，如此grad会自动更新为这两次的梯度之和。同理，若要累计k次的梯度信息，只要直接进行k次反向传播即可。

# example
x = torch.tensor([1.0], requires_grad=True)
y = torch.tensor([2.0])
z = x * y
z.backward(retain_graph=True)  
print(x.grad)  # tensor([2.])
z.backward(retain_graph=True)  
print(x.grad)  # tensor([4.])
z.backward(retain_graph=True)  
print(x.grad)  # tensor([6.])

当然，在更新网络模型时，我们其实不需要在反向传播函数中添加retain_graph=True。因为，我们每次都会重新设置计算图，但是由于模型的参数是不变的，因此之前计算的梯度值仍然存在。以下程序会累计4次的梯度信息，然后更新一次模型参数。

...
model = MLP()
optimizer.zero_grad()  # 清零grad
for step in range(1000):
    predict = model(x_input)
    loss = F.mse_loss(target, predict)
    loss.backward()  # 计算grad
    if step % 4 == 0:
        optimizer.step()  # 更新参数
        optimizer.zero_grad()  # 清零grad