神经网络里面经常要计算梯度,我们习惯用
loss.backward()
隐式地计算出所有参数的梯度。
在用神经网络求解PDE(如PINN)时,需要显式地计算梯度(导数)。这个时候 backward
就派不上用场了,那怎么显式地去计算呢?这就要用到 autograd
的 grad
方法。
理论
神经网络的求导和PDE里的求导一个重要的区别是:神经网络里 backward
是标量对张量求导,并且想得到同形张量;PDE是矢量对矢量求导,导数是矢量,三个矢量形状相同,何解?
- 神经网络中在求导时,所谓的标量就是损失函数的值,所有样本的loss平均下来就是一个标量。求导时对每个权重矩阵或偏置向量求导,从梯度下降的目的来看,就是要修正每个参数,所以得到的肯定是同形张量。
- PDE是矢量对矢量求导,这可能跟大多数人想的不一样,他们可能认为是标量对标量求导。是的,数学上确实是标量对标量求导。但我说的是用神经网络求解PDE这个情境。比如说我们在给定的区域内均匀采集了
n
n
n 个点,每个点都有一个
x
,
y
,
z
,
t
x,y,z,t
x,y,z,t。那你说,这
n
n
n 个点的集合,不就是4个矢量嘛,分别是
x
∈
R
n
,
y
∈
R
n
,
z
∈
R
n
,
t
∈
R
n
\boldsymbol{x}\in\mathbb{R}^n,\boldsymbol{y}\in\mathbb{R}^n,\boldsymbol{z}\in\mathbb{R}^n,\boldsymbol{t}\in\mathbb{R}^n
x∈Rn,y∈Rn,z∈Rn,t∈Rn,同时我们还可以算出每个点对应的
u
u
u。构成矢量
u
∈
R
n
\boldsymbol{u}\in\mathbb{R}^n
u∈Rn。那我们数学上的
∂
u
/
∂
x
\partial u/\partial x
∂u/∂x,到了代码里,因为有了
n
n
n 个点,可不就是矢量对矢量求导
∂
u
/
∂
x
\partial\boldsymbol{u}/\partial \boldsymbol{x}
∂u/∂x 了嘛。只不过这种矢量对矢量求导和一般的不同,这里不同维度的分量(不同点)之间互不干预,
u
1
u_1
u1 只跟
x
1
,
y
1
,
z
1
,
t
1
x_1,y_1,z_1,t_1
x1,y1,z1,t1 有关,跟其他量无关。而我们想得到的也是一个
n
n
n 维矢量:
( ∂ u ∂ x ) wanted = [ ∂ u 1 ∂ x 1 ∂ u 2 ∂ x 2 ⋯ ∂ u n ∂ x n ] T \left(\frac{\partial\boldsymbol{u}}{\partial\boldsymbol{x}}\right)_\text{wanted}=\begin{bmatrix} \dfrac{\partial u_1}{\partial x_1} & \dfrac{\partial u_2}{\partial x_2} & \cdots & \dfrac{\partial u_n}{\partial x_n} \end{bmatrix}^T (∂x∂u)wanted=[∂x1∂u1∂x2∂u2⋯∂xn∂un]T
这边之所以使用下标 wanted \text{wanted} 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 n n 维偏导数呢?
说来也简单,我们之前说了,不同维度之间互不干预,所以
J
J
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/∂x10⋮00∂u2/∂x2⋮0⋯⋯⋱⋯00⋮∂un/∂xn
11⋮1
求一阶偏导
想到了这一点,就可以看代码了。比如我们要求 y = x 1 2 sin x 2 y=x_1^2\sin x_2 y=x12sinx2 对两个自变量的偏导数。随机采集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个元素就可以。
这个元素是什么呢?就是 [ ∂ y 1 ∂ x 1 ∂ y 2 ∂ x 2 ] T = [ 2 x 1 sin x 2 x 1 2 cos x 2 ] T \begin{bmatrix} \dfrac{\partial y_1}{\partial x_1} & \dfrac{\partial y_2}{\partial x_2} \end{bmatrix}^T=\begin{bmatrix} 2x_1\sin x_2 & x_1^2\cos x_2 \end{bmatrix}^T [∂x1∂y1∂x2∂y2]T=[2x1sinx2x12cosx2]T。同时我们也要注意,PyTorch里二阶张量不同的维度是用不同的列表示的,所以这边的结果是一个行向量 [ 2 x 1 sin x 2 x 1 2 cos x 2 ] \begin{bmatrix} 2x_1\sin x_2 & x_1^2\cos x_2 \end{bmatrix} [2x1sinx2x12cosx2]。验证一下:
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
。我们让
∂
y
/
∂
x
1
=
2
x
1
sin
x
2
\partial y/\partial x_1=2x_1\sin x_2
∂y/∂x1=2x1sinx2 对
x
1
x_1
x1 和
x
2
x_2
x2 求导。
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求二阶导
这次求出来的应该是 [ 2 sin x 2 2 x 1 cos x 2 ] \begin{bmatrix} 2\sin x_2 & 2x_1\cos x_2 \end{bmatrix} [2sinx22x1cosx2],验证一下:
my2 = torch.hstack((
2 * torch.sin(x[:, [1]]),
2 * x[:, [0]] * torch.cos(x[:, [1]])
))
print(my2.equal(grad2))
True