Pytorch中backward()的思考記錄

在pytorch中，只能對標量使用backward。如果對向量進行backward，則會報錯：

import torch

x=torch.tensor([2,3,4],dtype=torch.float,requires_grad=True)
print(x)
y=x*2
print(y)
# z=y.mean()
# z.backward()
y.backward()
print(x.requires_grad)
print(x.grad)

以上程式碼執行：RuntimeError: grad can be implicitly created only for scalar outputs

但如果我們將y.backward()改為y.backward(torch.Tensor([1,1,1]))，則可以正確輸出x的梯度。

結合官方文件和知乎部落格https://zhuanlan.zhihu.com/p/27808095，進行以下直觀解釋：

我們有計算序列$m=(x_1=2,x_2=3),k=(x_1^2+3x_2,x_2^2+2x_1)$。為了方便表示，我們令$y_1=x_1^2+3x_2,y_2=x_2^2+2x_1$則$k=(y_1,y_2)$。我們在進行神經網路計算時，通常想要得到的梯度是$\frac{\partial y_1}{\partial x_1},\frac{\partial y_1}{\partial x_2},\frac{\partial y_2}{\partial x_1},\frac{\partial y_2}{\partial x_2}$，並以此進行反向傳播並優化。

但此時，我們的輸出$k$是一個向量，我們直接對$k$進行反向傳播：
$$\begin{gathered} \frac{\partial k}{\partial x_1}=\frac{\partial k}{\partial y_1}\cdot \frac{\partial y_1}{\partial x_1}+\frac{\partial k}{\partial y_2}\cdot \frac{\partial y_2}{\partial x_1} \\ \frac{\partial k}{\partial x_2}=\frac{\partial k}{\partial y_1}\cdot \frac{\partial y_1}{\partial x_2}+\frac{\partial k}{\partial y_2}\cdot \frac{\partial y_2}{\partial x_2} \end{gathered}$$
我們可以看到，上述直接對$k$進行反向傳播得到的結果並不是我們想要的。

寫成矩陣運算形式：
$$\left[\begin{array}{c}\frac{\partial k}{\partial x_1}\\ \frac{\partial k}{\partial x_2}\end{array}\right]=\left[\begin{array}{cc}\frac{\partial y_1}{\partial x_1}& \frac{\partial y_2}{\partial x_1}\\ \frac{\partial y_1}{\partial x_2}& \frac{\partial y_2}{\partial x_2}\end{array}\right]\left[\begin{array}{c}\frac{\partial k}{\partial y_1}\\ \frac{\partial k}{\partial y_2}\end{array}\right]$$
可以看到，矩陣運算形式中前面$2\times 2$的矩陣中的結果正是我們想要的。我們稱之為雅可比(Jacobian)矩陣(嚴格地說，是雅可比矩陣的轉置)。

那麼，當我們令$\left[\begin{array}{c}\frac{\partial k}{\partial y_1}\\ \frac{\partial k}{\partial y_2}\end{array}\right]=\left[\begin{array}{c}1\\ 0\end{array}\right]$時，即可得到雅可比矩陣中的第一列，令$\left[\begin{array}{c}\frac{\partial k}{\partial y_1}\\ \frac{\partial k}{\partial y_2}\end{array}\right]=\left[\begin{array}{c}0\\ 1\end{array}\right]$時，即可得到雅可比矩陣的第二列。以此我們就可以得到想要的梯度值。

這正是pytorch官方文件中所說的思想。當我們對向量進行反向傳播時，通過在backward()中新增一個向量$v$，就可以分別得到原向量中每一項乘向量$v$中係數後對應的梯度值。

一些需要注意的點：

1. 在backward()中新增的向量$v$的size要與進行反向傳播的向量的size相同。
2. 當我們需要分別求反向傳播向量中每一項的梯度時，可能需要分多次分別進行求導。此時我們要注意backward()裡面另外的一個引數retain_variables=True，這個引數預設是False，也就是反向傳播之後這個計算圖的記憶體會被釋放，這樣就沒辦法進行第二次反向傳播了，所以我們需要設定為True，因為這裡我們需要進行兩次反向傳播求得Jacobian矩陣。