pytorch教程之自动求导机制(AUTOGRAD)-从梯度和Jacobian矩阵讲起

---

1. 梯度和Jacobian矩阵

设$f(x)\in R^1$是关于向量$x\in R^n$的函数，则它关于$x$的导数定义为：
$$\frac{df(x)}{dx}:=\left[\frac{\partial f(x)}{\partial x_i}\right]\in R^n\text{\hspace{1em}(1-1)}$$
函数$f(x)\in R^1$关于向量$x\in R^n$的导数是一个列向量，称之为$f(x)$关于$x$的梯度。
$$\frac{df(x{)}^T}{dx}:={\left(\frac{df(x)}{dx}\right)}^T={\left[\frac{\partial f(x)}{\partial x_i}\right]}^T\in R^{1\times n}\text{\hspace{1em}(1-2)}$$
如果$f(x)\in R^M$是关于向量$x\in R^n$的函数向量，则$f(x)$关于$x$的导数定义为：
$$\frac{df(x)}{dx}:=\frac{df(x)}{d{x}^T}=\left[\frac{\partial f(x)}{\partial x_1},\frac{\partial f(x)}{\partial x_2},\cdots ,\frac{\partial f(x)}{\partial x_n}\right]\in R^{m\times n}\text{\hspace{1em}(1-3)}$$
称上述矩阵为Jacobian矩阵。
一些常用推论：

1. 假设$v,x\in R^n$:
   $$\frac{d}{dx}(v^{T}x)=\frac{d}{dx}(x^{T}v)=v\text{\hspace{1em}(1-4)}$$
2. 假设$y\in R^1$,$z\in R^m$,$x\in R^n$,$z=g(x)$,y=f(z)：
   $$\frac{dy}{dx}={\left(\frac{dz}{dx}\right)}^T\frac{dy}{dz}\text{\hspace{1em}(1-5)}$$
   可以从向量矩阵的维度适配上去理解和记忆，因为$\frac{dy}{dx}\in R^n$,$\frac{dy}{dz}\in R^m$,$\frac{dz}{dx}\in R^{m\times n}$，所以必须有上述的公式才能适配。
3. 假设$y\in R^k$,$z\in R^m$,$x\in R^1$,$z=g(x)$,y=f(z)：
   $$\frac{dy}{dx}=\frac{dy}{dz}\frac{dz}{dx}\text{\hspace{1em}(1-6)}$$
4. 假设$y\in R^k$,$z\in R^m$,$x\in R^n$,$z=g(x)$,y=f(z)：
   $$\frac{dy}{dx}=\frac{dy}{dz}\frac{dz}{dx}\text{\hspace{1em}(1-7)}$$

2. pytorch求变量导数的过程

在pytorch和TensorFlow中，是不支持张量对张量的求导。这不是因为数学上没法求，而是因为工程实现上比较麻烦。因为向量对向量求导是个矩阵，二阶张量（矩阵）对二阶张量（矩阵）求导得到一个四阶张量，这样很容易会产生阶数爆炸。所以pytorch和TensorFlow（猜测其他深度学习框架也是这样）对外的接口干脆不支持张量对张量求导。如果遇到张量对张量求导的情况，例如向量对向量求导的情况，需要对因变量乘以一个维度一样的向量，转换为标量对向量的求导，这样可以大大减少计算量（具体见后文）。并且，因为pytorch和TensorFlow是为了机器学习/深度学习模型设计的，机器学习/深度模型的求导基本上都是损失函数(标量)对参数的求导，很少直接用到向量对向量求导，因此上述过程是有实际意义和需求的。

假设有一个三维tensor $x=[x_1,x_2,x_3{]}^T=[1,2,3{]}^T$,另一个三维tensor y:
$$y=f(x)=\left[\begin{array}{c}{x_1}^3+2{x_2}^2+3x_3\\ 3x_1+2{x_2}^2+{x_3}^3\\ 2x_1+{x_2}^3+3{x_3}^2\end{array}\right]\text{\hspace{1em}(2-1)}$$
那么在计算y相对于x的导数时，
$$\frac{dy}{dx}=\left[\begin{array}{cccc}& 3{x_1}^2,& 4x_2,& 3\\ & 3,& 4x_2,& 3{x_3}^2\\ & 2,& 3{x_2}^2,& 6x_3\end{array}\right]\text{\hspace{1em}(2-2)}$$
在pytorch中实际计算时，不能直接用y对x求导，需要先用一个向量$w$左乘y，再转置。例如，$w^T=[3,2,1]$。因此，pytorch算的其实是：
$$\frac{d{y}^T}{dx}w={\left(w^T\frac{dy}{dx}\right)}^T=\left[\begin{array}{c}17\\ 52\\ 81\end{array}\right]\text{\hspace{1em}(2-3)}$$
$w$可以理解为是对$[\frac{\partial y_1}{\partial x},\frac{\partial y_2}{\partial x},\frac{\partial y_3}{\partial x}{]}^T$的权重参数。因此我们得到的是y的各个分量的导数的加权求和。

代码如下：

import torch
x1=torch.tensor(1, requires_grad=True, dtype = torch.float)
x2=torch.tensor(2, requires_grad=True, dtype = torch.float)
x3=torch.tensor(3, requires_grad=True, dtype = torch.float)
y=torch.randn(3)
y[0]=x1**3+2*x2**2+3*x3
y[1]=3*x1+2*x2**2+x3**3
y[2]=2*x1+x2**3+3*x3**2
v=torch.tensor([3,2,1],dtype=torch.float)
y.backward(v)
print(x1.grad)
print(x2.grad)
print(x3.grad)

利用链式求导的原理来理解，可以理解为$w$是（远方）某个标量对$y$的导数。pytorch之所以要这么设计，是因为在机器学习/深度学习模型中，求导的最终目的一般是为了让损失函数最小。损失函数一般都是一个标量，因此无论链式求导的过程多么复杂，中间过程也许有很多向量对向量求导的子过程，但是最开始一定会有一个标量（损失函数）对向量的求导过程，这个导数就是前面的$w$。

下面看一个带两个隐藏层的神经网络解决线性回归问题的例子，来进一步说明这点。
为了简单起见，考虑batch_size=1的情况。设输入数据为$x=[x_1,x_2{]}^T$,输入层到第一个隐藏层的权重矩阵为
$$W=\left[\begin{array}{c}{w}_1^T\\ w_2^T\end{array}\right]=\left[\begin{array}{c}{w}_{11},w_{12}\\ w_{21},w_{22}\end{array}\right]\text{\hspace{1em}(2-4)}$$
第一个隐藏层的值为$z=[z_1,z_2{]}^T$,
第一个隐藏层到第二个隐藏层的权重矩阵为
$$U=\left[\begin{array}{c}{u}_1^T\\ u_2^T\end{array}\right]=\left[\begin{array}{c}{u}_{11},u_{12}\\ u_{21},u_{22}\end{array}\right]\text{\hspace{1em}(2-5)}$$
第二个隐藏层的值为$s=[s_1,s_2{]}^T$,
输出层的值为$y$,隐藏层到输出层的权重参数为$v=[v_1,v_2{]}^T$。则有：
$$z=\left[\begin{array}{c}{z}_1\\ z_2\end{array}\right]=\left[\begin{array}{c}{w}_{11},w_{12}\\ w_{21},w_{22}\end{array}\right]\left[\begin{array}{c}{x}_1\\ x_2\end{array}\right]=\left[\begin{array}{c}{w}_{11}{x}_1+w_{12}{x}_2\\ w_{21}{x}_1+w_{22}{x}_2\end{array}\right]\text{\hspace{1em}(2-6)}$$

$$\begin{array}{rl}s& =\left[\begin{array}{c}{s}_1\\ s_2\end{array}\right]=\left[\begin{array}{c}{u}_{11},u_{12}\\ u_{21},u_{22}\end{array}\right]\left[\begin{array}{c}{z}_1\\ z_2\end{array}\right]\\ & =\left[\begin{array}{c}{u}_{11}(w_{11}{x}_1+w_{12}{x}_2)+u_{12}(w_{21}{x}_1+w_{22}{x}_2)\\ u_{21}(w_{11}{x}_1+w_{12}{x}_2)+u_{22}(w_{21}{x}_1+w_{22}{x}_2)\end{array}\right]\end{array}\text{\hspace{1em}(2-7)}$$

$$\begin{array}{rl}y& =[v_1,v_2]\left[\begin{array}{c}{s}_1\\ s_2\end{array}\right]\\ & =(v_1{u}_{11}{x}_1+v_2{u}_{21}{x}_1)w_{11}\\ & +(v_1{u}_{11}{x}_2+v_2{u}_{21}{x}_2)w_{12}\\ & +(v_1{u}_{12}{x}_1+v_2{u}_{22}{x}_1)w_{21}\\ & +(v_1{u}_{12}{x}_2+v_2{u}_{22}{x}_2)w_{22}\end{array}\text{\hspace{1em}(2-8)}$$
损失函数为$L=(y-\hat{y}{)}^2/2$
则损失函数关于权重参数$w_1$的导数为：
$$\begin{array}{rl}\frac{dL}{d{w}_1}& =(y-\hat{y})\frac{dy}{dx}\\ & =(y-\hat{y})\frac{d{s}^T}{dx}\frac{dy}{ds}\\ & =(y-\hat{y})\frac{d{z}^T}{dx}\frac{d{s}^T}{dz}\frac{dy}{ds}\\ & =(y-\hat{y})\left[\begin{array}{c}{x}_1,0\\ x_2,0\end{array}\right]\left[\begin{array}{c}{u}_{11},u_{21}\\ u_{12},u_{22}\end{array}\right]\left[\begin{array}{c}{v}_1\\ v_2\end{array}\right]\\ & =(y-\hat{y})\left[\begin{array}{c}{v}_1{x}_1{u}_{11}+v_2{x}_1{u}_{21}\\ v_1{x}_2{u}_{11}+v_2{x}_2{u}_{21}\end{array}\right]\end{array}\text{\hspace{1em}(2-9)}$$
可以验证$(2-9)$和前面$(2-8)$中直接求得的导数值是一样的。
这里发现了一个小彩蛋：
假设在pytorch的底层实现中，如果从左往右计算，则需要进行进行大量的矩阵乘法。如果有n个$2\times 2$的方阵相乘，那么需要进行$4\times (n-1)$次内积。如果从又往左计算，只需要进行$2\times n$次内积。