在PyTorch中计算导数

神经网络里面经常要计算梯度，我们习惯用 loss.backward() 隐式地计算出所有参数的梯度。

在用神经网络求解PDE（如PINN）时，需要显式地计算梯度（导数）。这个时候 backward 就派不上用场了，那怎么显式地去计算呢？这就要用到 autograd 的 grad 方法。

理论

神经网络的求导和PDE里的求导一个重要的区别是：神经网络里 backward 是标量对张量求导，并且想得到同形张量；PDE是矢量对矢量求导，导数是矢量，三个矢量形状相同，何解？

• 神经网络中在求导时，所谓的标量就是损失函数的值，所有样本的loss平均下来就是一个标量。求导时对每个权重矩阵或偏置向量求导，从梯度下降的目的来看，就是要修正每个参数，所以得到的肯定是同形张量。
• PDE是矢量对矢量求导，这可能跟大多数人想的不一样，他们可能认为是标量对标量求导。是的，数学上确实是标量对标量求导。但我说的是用神经网络求解PDE这个情境。比如说我们在给定的区域内均匀采集了 $n$ 个点，每个点都有一个 $x,y,z,t$。那你说，这 $n$ 个点的集合，不就是4个矢量嘛，分别是 $\mathbf{x}\in {\mathbb{R}}^n,\mathbf{y}\in {\mathbb{R}}^n,\mathbf{z}\in {\mathbb{R}}^n,\mathbf{t}\in {\mathbb{R}}^n$，同时我们还可以算出每个点对应的 $u$。构成矢量 $\mathbf{u}\in {\mathbb{R}}^n$。那我们数学上的 $\partial u/\partial x$，到了代码里，因为有了 $n$ 个点，可不就是矢量对矢量求导 $\partial \mathbf{u}/\partial \mathbf{x}$ 了嘛。只不过这种矢量对矢量求导和一般的不同，这里不同维度的分量（不同点）之间互不干预，$u_1$ 只跟 $x_1,y_1,z_1,t_1$ 有关，跟其他量无关。而我们想得到的也是一个 $n$ 维矢量：
  $${\left(\frac{\partial \mathbf{u}}{\partial \mathbf{x}}\right)}_{\text{wanted}}={\left[\begin{array}{cccc}\frac{\partial u_1}{\partial x_1}& \frac{\partial u_2}{\partial x_2}& \cdots & \frac{\partial u_n}{\partial x_n}\end{array}\right]}^T$$
  这边之所以使用下标 $\text{wanted}$，就是表明这不是数学上应该得到的计算结果。数学上，矢量对矢量求导得到一个Jacobian矩阵：
  J = ∂ u ∂ x = [ ∂ u 1 / ∂ x 1 ∂ u 1 / ∂ x 2 ⋯ ∂ u 1 / ∂ x n ∂ u 2 / ∂ x 1 ∂ u 2 / ∂ x 2 ⋯ ∂ u 2 / ∂ x n ⋮ ⋮ ⋱ ⋮ ∂ u n / ∂ x 1 ∂ u n / ∂ x 2 ⋯ ∂ u n / ∂ x n ] J=\frac{\partial\boldsymbol{u}}{\partial\boldsymbol{x}}=\begin{bmatrix} \partial u_1/\partial x_1 & \partial u_1/\partial x_2 & \cdots & \partial u_1/\partial x_n\\ \partial u_2/\partial x_1 & \partial u_2/\partial x_2 & \cdots & \partial u_2/\partial x_n\\ \vdots & \vdots & \ddots & \vdots \\ \partial u_n/\partial x_1 & \partial u_n/\partial x_2 & \cdots & \partial u_n/\partial x_n \end{bmatrix} J=∂x∂u​= ​∂u1​/∂x1​∂u2​/∂x1​⋮∂un​/∂x1​​∂u1​/∂x2​∂u2​/∂x2​⋮∂un​/∂x2​​⋯⋯⋱⋯​∂u1​/∂xn​∂u2​/∂xn​⋮∂un​/∂xn​​ ​
  那怎样才能从Jacobian矩阵得到我们想要的 $n$ 维偏导数呢？

说来也简单，我们之前说了，不同维度之间互不干预，所以 $J$ 一定是一个对角阵，非对角元素均为零。现在问题就变成了怎么把对角阵的对角元素提取出来，一个简单的想法是右乘一个全1向量：
( ∂ u ∂ x ) wanted = [ ∂ u 1 / ∂ x 1 0 ⋯ 0 0 ∂ u 2 / ∂ x 2 ⋯ 0 ⋮ ⋮ ⋱ ⋮ 0 0 ⋯ ∂ u n / ∂ x n ] [ 1 1 ⋮ 1 ] \left(\frac{\partial\boldsymbol{u}}{\partial\boldsymbol{x}}\right)_\text{wanted}=\begin{bmatrix} \partial u_1/\partial x_1 & 0 & \cdots & 0\\ 0 & \partial u_2/\partial x_2 & \cdots & 0\\ \vdots & \vdots & \ddots & \vdots \\ 0 & 0 & \cdots & \partial u_n/\partial x_n \end{bmatrix}\begin{bmatrix} 1 \\ 1\\ \vdots \\ 1 \end{bmatrix} (∂x∂u​)wanted​= ​∂u1​/∂x1​0⋮0​0∂u2​/∂x2​⋮0​⋯⋯⋱⋯​00⋮∂un​/∂xn​​ ​ ​11⋮1​ ​

求一阶偏导

想到了这一点，就可以看代码了。比如我们要求 $y=x_1^2\sin x_2$ 对两个自变量的偏导数。随机采集100个点。

from torch import autograd
import torch

n = 100
x = torch.rand(n, 2, requires_grad=True)
y = x[:, 0] ** 2 * torch.sin(x[:, 1])
grad = autograd.grad(
    outputs=y, inputs=x,
    grad_outputs=torch.ones_like(y),
)[0]

这里面用到的一个重要方法是 torch.autograd.grad。我们用 outputs 指定因变量，用 inputs 指定自变量，用 grad_outputs 指定右乘的向量。该函数返回一个元组，通常会是单元素元组，我们取出第0个元素就可以。

这个元素是什么呢？就是 ${\left[\begin{array}{cc}\frac{\partial y_1}{\partial x_1}& \frac{\partial y_2}{\partial x_2}\end{array}\right]}^T={\left[\begin{array}{cc}2x_1\sin x_2& x_1^2\cos x_2\end{array}\right]}^T$。同时我们也要注意，PyTorch里二阶张量不同的维度是用不同的列表示的，所以这边的结果是一个行向量 $\left[\begin{array}{cc}2x_1\sin x_2& x_1^2\cos x_2\end{array}\right]$。验证一下：

my = torch.hstack((
    2 * x[:, [0]] * torch.sin(x[:, [1]]),
    x[:, [0]] ** 2 * torch.cos(x[:, [1]])))
print(my.size())
print(my.equal(grad))

torch.Size([100, 2])
True

求二阶偏导

求二阶导要利用一阶导的结果，并且要求我们在求一阶导时必须指定 create_graph 为 True。我们让 $\partial y/\partial x_1=2x_1\sin x_2$ 对 $x_1$ 和 $x_2$ 求导。

grad = autograd.grad(
    outputs=y, inputs=x,
    grad_outputs=torch.ones_like(y),
    create_graph=True
)[0]  # 一阶导
grad2 = autograd.grad(
    outputs=grad[:, 0], inputs=x,
    grad_outputs=torch.ones_like(grad[:, 0])
)[0]  # 对x_1求完一阶导后，再对x_1和x_2求二阶导

这次求出来的应该是 $\left[\begin{array}{cc}2\sin x_2& 2x_1\cos x_2\end{array}\right]$，验证一下：

my2 = torch.hstack((
    2 * torch.sin(x[:, [1]]),
    2 * x[:, [0]] * torch.cos(x[:, [1]])
))
print(my2.equal(grad2))

True