Pytorch中的自动微分 | Fructose’s Notion Blog

type

Post

status

Published

date

Feb 16, 2023

slug

summary

深度学习框架通过自动计算导数，即自动微分 （automatic differentiation）来加快求导。实际中，根据设计好的模型，系统会构建一个计算图 （computational graph），来跟踪计算是哪些数据通过哪些操作组合起来产生输出。自动微分使系统能够随后反向传播梯度。这里，反向传播 （backpropagate）意味着跟踪整个计算图，填充关于每个参数的偏导数。

一元函数求导

假设我们需要对以下函数求导

通过数学，我们可以很简单地计算出它的导函数

但是每次都需要手动计算微分会非常麻烦，Pytorch中提供了一个工具自动微分可以很简单地解决这个问题，具体步骤如下：


import torch

def f(x):
    return torch.exp(x) * (x + torch.sin(x))

x = torch.tensor(2., requires_grad=True)
y = f(x)
y.backward()
print(x, y, x.grad)

# 输出
# (tensor(2., requires_grad=True),
#  tensor(21.4970, grad_fn=<MulBackward0>),
#  tensor(25.8111))

上述代码计算了x=2时的函数值和一阶导数值，此外，我们还可以验证它与我们自己计算的值是否相等：


def f_x(x):
    return torch.exp(x) * (x + torch.sin(x) + 1 + torch.cos(x))

with torch.no_grad():
    print(abs(x.grad - f_x(x)))

# 输出
# tensor(1.9073e-06)

从结果可以看出，两者的差距非常小，基本上可以看作相等。

连续求导

在进行反向传播的时候输入一个参数，即可实现连续求导，将自变量变为一个或者多个。以上面的函数为例：


x = torch.arange(-5., 5., 0.01, requires_grad=True)
y = f(x)
y.backward(torch.ones_like(x))
x.grad.shape

# 输出
# torch.Size([1000])

这样，就可以绘制其函数与其一阶导函数图像，这里我们还画出我们的解析解：


with torch.no_grad():
    y1 = f_x(x)
fig = plt.figure(figsize=(6, 4))
ax = fig.add_subplot(111)
ax.plot(x.detach().numpy(), y.detach().numpy(), label="f(x)")
ax.plot(x.detach().numpy(), x.grad.detach().numpy(), linestyle="solid", label="autograd")
ax.plot(x.detach().numpy(), y1.numpy(), linestyle="--", label="grad")
ax.set_title("f(x) and grad")
ax.set_xlabel("x")
ax.set_ylabel("f(x) and grad")
ax.legend()
fig.savefig("autograd.png", dpi=300)

画出的图片如下图所示：

从图中可以看出，绿色的虚线和黄色的实线重合，这表示自动微分的结果，和我们手动求导的结果相同。

一元函数的高阶导数

根据自动微分的原理，我们可以使用以下函数来求n阶导数


def nth_derivative(f, wrt, n):
    for i in range(n):
        if not f.requires_grad:
            return torch.zeros_like(wrt)
        grads = grad(f, wrt, create_graph=True)[0]
        f = grads.sum()
    return grads

其中f表示因变量，wrt表示自变量，n表示导数的阶数，需要大于0.

使用这个函数来计算上述函数的二阶导数的例子如下：


x = torch.tensor(3., requires_grad=True)
y = f(x)
print(nth_derivative(y, x, 2))

# 结果
# tensor(60.6586, grad_fn=<AddBackward0>)

矩阵梯度

对于矩阵（或向量）求梯度也可以使用自动微分：


x = torch.arange(4., requires_grad=True)
y = 2 * torch.dot(x, x)
y.backward()
x.grad

# 输出
# tensor([ 0.,  4.,  8., 12.])

💡

对于矩阵或者向量求梯度时，要保证最后的函数值是一个一维的数，如果不是，可以对结果进行sum()运算后再使用反向传播计算梯度。

扩展说明

在进行自动微分时的函数也可以是由Python的控制流决定的，例如：


def f(a):
    b = a * 2
    while b.norm() < 1000:
        b = b * 2
    if b.sum() > 0:
        c = b
    else:
        c = 100 * b
    return c

对于这种函数也可以使用自动微分求梯度。

test_autograd.ipynb

32.0KB

参考文章

2.5. 自动微分 — 动手学深度学习 2.0.0 documentation

https://zh.d2l.ai/chapter_preliminaries/autograd.html

4. Autograd (4)：PyTorch 高阶自动导数 — ajz34 文档

https://ajz34.readthedocs.io/zh_CN/latest/ML_Notes/Autograd_Series/Autograd_Hess.html#stackoverflow-qa