不过,在大部分场景中想计算的二阶导数是另一种形式:import torch
x = ...
loss = model(x)
w1_grad = torch.autograd.grad(outputs=loss, inputs=model.w1, create_graph=True)
w1_grad.backward()
它会报错RuntimeError: grad can be implicitly created only for scalar outputs,即 y 不是一个标量。 如果我们给y.backward()传入任意一个形状和 y 相同的 tensor 作为gradient参数,如y.backward(gradient=torch.tensor([2, 3])),这份代码就可以正常运行。import torch
x = torch.ones(2).requires_grad_()
y = x * 2
y.backward()
这背后的原因是,反向传播只能应用于一个单输出(即输出为标量)的函数。当我们设置# 省略了定义 x 和 y 的代码
2, 3])) > y.backward(gradient=torch.tensor([
4., 6.]) > x.gradtensor([
gradient
为 [2, 3] 时,PyTorch 会将 gradient 和 y 做内积,得到一个标量
t = 2 * y[0] + 3 * y[1]
, 将标量 t 当作函数的输出并运行反向传播,因此 x.grad 的值是
dt/dx
,根据标量对矩阵的导数的定义容易得到dt/dy = [dt/d(y[0]), dt/d(y[1])] = [2, 3],根据链式法则,dt/dx = dt/dy * dy/dx = [2, 3] * 2 = [4, 6],所以 x.grad 的值会是 [4, 6]。
也就是说,对于任意一个 N 输出 M 输入的函数Y = f(X)(Y 是一个长度为 N 的向量),我们可以将 Y 和 v 做内积,得到一个标量 t 和一个单输出、M 输入的函数 g 满足t = g(X) = v * Y = v * f(X)
,再通过反向传播求出函数 g 的梯度。
根据链式法则,函数 g 的梯度dt/dX等于dt/dY * dY/dX,而显然dt/dY就是向量 v,也就是 loss.backward() 函数中的 gradient
参数,dY/dX是函数f的导数(形状为 N*M 的 Jacobian 矩阵)。也就是说,以一个长度为 N 的向量 v 和一个 N 输出 M 输入的函数 f 为对象运行反向传播,可以求出向量 v 和函数 f 的导数的乘积,也就是 Vector-Jacobian Product,简称为 VJP。如果想用反向传播求出整个 Jacobian 矩阵,那么可以将 v 设置为第 i 个元素为 1,其它元素均为 0 的向量,这样就可以求出 Jacobian 矩阵的第 i 行,重复 N 次就可以求出整个矩阵。
同样根据链式法则,dY/dt = dY/dX * dX/dt,dY/dX是 f 的导数(Jacobian 矩阵) ,dX/dt就是向量 v。因此,以一个长度为 M 的向量 v 和一个 N 输出M 输入的函数 f 为对象运行前向自动微分,可以求出函数 f 的导数和向量 v 的乘积,这个积可以被简称为 JVP(Jacobian-Vector Product)。
类似地,如果想用前向自动微分求出整个 Jacobian 矩阵,也是可以将 v 设置为第 i 个元素为 1,其它元素均为 0 的向量来求出 Jacobian 矩阵的第 i 列,重复 M 次就可以求出整个矩阵。
由此可见,如果是想求出整个 Jacobian 矩阵,理论上来说,对于 M > N,也就是输入数量大于输出数量的函数,使用反向传播会更高效,否则使用前向自动微分会更高效。 不过,在实际中,因为前向自动微分是根据输入以及输入的导数一次性算出输出和输出的导数,对 cache 更友好,所以即使输入数量和输出数量相等,前向自动微分的速度也会比反向传播更快。 PyTorch 里也已经有了实验性的 Forward Mode 自动微分的 API,https://pytorch.org/docs/master/autograd.html#forward-mode-automatic-differentiation 。自动微分的经典书籍 《Evaluating Derivatives》里甚至总结了这样一条规则:不要做超过一次反向微分。 再回过头来看文章最初的 Stackoverflow 回答,我们现在已有的知识已经可以分析出它到底是在求什么东西了。它在每一次调用grad函数后,将结果grads进行了求和,再对求和结果进行下一次后向微分,相当于通过乘以一个全 1 的向量,将多输出的函数变成了单输出,再进行后向微分。 因此 n=2 时它求的是 v 为全 1 向量时的 HVP,n=3 时是以全 1 向量和 n=2 的结果为对象求 VJP,依次类推。这些东西可能在某些场景下也有用,但显然不能说是 general 的 “高阶导数”。 让我们验证一下我们的分析:Forward over reverse
3.66 ms ± 89.7 µs per loop (mean ± std. dev. of 3 runs, 10 loops each)
Reverse over forward
7.43 ms ± 4.22 ms per loop (mean ± std. dev. of 3 runs, 10 loops each)
Reverse over reverse
11.1 ms ± 6.9 ms per loop (mean ± std. dev. of 3 runs, 10 loops each)
Naive full Hessian materialization
55 ms ± 476 µs per loop (mean ± std. dev. of 3 runs, 10 loops each)
这样一段代码,求的是 当 x0=1, x1=2 时的 2 阶“导数”,通过手算很容易得到 loss 的 Hessian 矩阵是 ,然而这段代码打印的结果是 [216., 264.] ,和上述的分析一致。 目前我们已经知道了 HVP 的求法,那么和求完整 Jacobian 矩阵时一样,用不同的 v 去计算 HVP,重复多次之后就可以得到完整的 Hessian 矩阵。但实际应用中,Hessian 矩阵的大小很可能非常非常大,例如对神经网络训练来说,Hessian 矩阵的大小是 N*N,N 是神经网络中的所有权重的元素数量之和,在 10^9 量级,要完整求出这么大的矩阵是不现实的。 所以通常都会想办法绕过 Hessian 矩阵的计算,例如用性质类似的矩阵代替 Hessian 矩阵,或者通过变换用 HVP 代替 Hessian 矩阵的功能(拟牛顿法、共轭梯度法)。import torch
from torch.autograd import grad
# 注:这个代码是反例,不要直接拿去用
def nth_derivative(f, wrt, n):
for i in range(n):
grads = grad(f, wrt, create_graph=True)[0]
f = grads.sum()
return grads
x = torch.arange(1, 3, dtype=torch.float, requires_grad=True)
x = x.reshape(2, 1)
x = torch.cat((x[0] + x[1], x[1]), 0)
y = x ** 4
loss = y.sum()
print(nth_derivative(f=loss, wrt=x0, n=2))
Taylar Mode 求高阶导数相当于将前向自动微分涉及的对象从实数扩展为了 k 阶泰勒多项式,Jax 的相应 API 是https://jax.readthedocs.io/en/latest/jax.experimental.jet.html。它在深度学习框架界是一个很新的东西,也可以用在反向自动微分里,待更新......
(注:这篇文章假设读者了解梯度和 Jacobian 矩阵的定义,如果不了解可以阅读维基百科。文章一些地方没有完全严谨,比如一些向量可能该加上转置符号。LaTeX 和 Markdown 混打比较麻烦所以只有复杂符号用了 LaTeX。)
其他人都在看本文分享自微信公众号 - OneFlow(OneFlowTechnology)。
如有侵权,请联系 support@oschina.cn 删除。
本文参与“OSC源创计划”,欢迎正在阅读的你也加入,一起分享。
|