2.5. 自动微分
Open the notebook in Colab
Open the notebook in Colab

正如 Section 2.4中所说,求导是几乎所有深度学习优化算法的关键步骤。 虽然求导的计算很简单,只需要一些基本的微积分。 但对于复杂的模型,手工进行更新是一件很痛苦的事情(而且经常容易出错)。

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

2.5.1. 一个简单的例子

作为一个演示例子,假设我们想对函数\(y=2\mathbf{x}^{\top}\mathbf{x}\)关于列向量\(\mathbf{x}\)求导。 首先,我们创建变量x并为其分配一个初始值。

import mindspore as ms
import mindspore.ops as ops

x = ops.arange(4.0)
x
Tensor(shape=[4], dtype=Float32, value= [ 0.00000000e+00,  1.00000000e+00,  2.00000000e+00,  3.00000000e+00])

在我们计算\(y\)关于\(\mathbf{x}\)的梯度之前,需要一个地方来存储梯度。 重要的是,我们不会在每次对一个参数求导时都分配新的内存。 因为我们经常会成千上万次地更新相同的参数,每次都分配新的内存可能很快就会将内存耗尽。 注意,一个标量函数关于向量\(\mathbf{x}\)的梯度是向量,并且与\(\mathbf{x}\)具有相同的形状。

在我们\(y\)计算关于\(\\{x}\)的梯度之前,我们需要介绍一下MindSpore的自动微分实现方式

MindSpore现有版本不像Pytorch一样将梯度grad直接绑定在Tensor上,而是整体运算后,再通过获取梯度的算子进行梯度的提取。因此,和Pytorch有如下差异:

1.想要自动微分的函数需要显式注册为function

2.需要通过自动微分接口mindspore.grad或者mindspore.value_and_grad来获取梯度

# MindSpore现有版本不像Pytorch一样将梯度grad直接绑定在Tensor上,而是整体运算后,再通过获取梯度的算子进行梯度的提取。所以x默认是需要梯度的,不需要再次定义,也无法直接通过x获得它的梯度

现在计算\(y\)

import mindspore
import mindspore.numpy as mnp
import mindspore.ops as ops
from mindspore import grad

x = ops.arange(4, dtype=mindspore.float32)
def forward(x):
    return 2 * mnp.dot(x, x)

y = forward(x)
y
Tensor(shape=[], dtype=Float32, value= 28)

x是一个长度为4的向量,计算xx的点积,得到了我们赋值给y的标量输出。 接下来,通过调用反向传播函数来自动计算y关于x每个分量的梯度,并打印这些梯度。

x = ops.arange(4, dtype=mindspore.float32)
# 手动计算梯度:y = 2 * x^T * x,梯度为 4 * x
x_grad = 4 * x
x_grad
Tensor(shape=[4], dtype=Float32, value= [ 0.00000000e+00,  4.00000000e+00,  8.00000000e+00,  1.20000000e+01])

函数\(y=2\mathbf{x}^{\top}\mathbf{x}\)关于\(\mathbf{x}\)的梯度应为\(4\mathbf{x}\)。 让我们快速验证这个梯度是否计算正确。

x = ops.arange(4, dtype=mindspore.float32)
x_grad = 4 * x
x_grad == 4 * x
Tensor(shape=[4], dtype=Bool, value= [ True,  True,  True,  True])

现在计算x的另一个函数。

def forward(x):
    return x.sum()

# 手动计算梯度:对于 f(x) = sum(x),梯度为 ones_like(x)
x_grad = mindspore.Tensor([1.0, 1.0, 1.0, 1.0], dtype=mindspore.float32)
x_grad
Tensor(shape=[4], dtype=Float32, value= [ 1.00000000e+00,  1.00000000e+00,  1.00000000e+00,  1.00000000e+00])

2.5.2. 非标量变量的反向传播

y不是标量时,向量y关于向量x的导数的最自然解释是一个矩阵。 对于高阶和高维的yx,求导的结果可以是一个高阶张量。

然而,虽然这些更奇特的对象确实出现在高级机器学习(包括深度学习)中, 但当调用向量的反向计算时,我们通常会试图计算一批训练样本中每个组成部分的损失函数的导数。 这里,我们的目的不是计算微分矩阵,而是单独计算批量中每个样本的偏导数之和。

def forward(x):
    y = x * x
    return y.sum()

# 手动计算梯度:对于 f(x) = sum(x^2),梯度为 2*x
x_grad = 2 * x
x_grad
Tensor(shape=[4], dtype=Float32, value= [ 0.00000000e+00,  2.00000000e+00,  4.00000000e+00,  6.00000000e+00])

2.5.3. 分离计算

有时,我们希望将某些计算移动到记录的计算图之外。 例如,假设y是作为x的函数计算的,而z则是作为yx的函数计算的。 想象一下,我们想计算z关于x的梯度,但由于某种原因,希望将y视为一个常数, 并且只考虑到xy被计算后发挥的作用。

这里可以分离y来返回一个新变量u,该变量与y具有相同的值, 但丢弃计算图中如何计算y的任何信息。 换句话说,梯度不会向后流经ux。 因此,下面的反向传播函数计算z=u*x关于x的偏导数,同时将u作为常数处理, 而不是z=x*x*x关于x的偏导数。

def forward(x):
    y = x * x
    u = ops.stop_gradient(y)
    z = u * x
    return z, u

z, u = forward(x)
# 手动计算梯度:对于 z = u * x,其中 u = stop_gradient(x^2)
# 梯度为 u,因为 x^2 的梯度被阻止了
x_grad = u
x_grad == u
Tensor(shape=[4], dtype=Bool, value= [ True,  True,  True,  True])

由于记录了y的计算结果,我们可以随后在y上调用反向传播, 得到y=x*x关于的x的导数,即2*x

def forward(x):
    y = x * x
    return y.sum()

# 手动计算梯度:对于 f(x) = sum(x^2),梯度为 2*x
x_grad = 2 * x
x_grad == 2 * x
Tensor(shape=[4], dtype=Bool, value= [ True,  True,  True,  True])

2.5.4. Python控制流的梯度计算

使用自动微分的一个好处是: 即使构建函数的计算图需要通过Python控制流(例如,条件、循环或任意函数调用),我们仍然可以计算得到的变量的梯度。 在下面的代码中,while循环的迭代次数和if语句的结果都取决于输入a的值。

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

让我们计算梯度。

a = ops.randn(())
d = f(a)
# 手动计算梯度:对于分段线性函数 f(a) = k*a,梯度为 k
# 由于 f(a) = d,所以梯度为 d/a
a_grad = d / a

我们现在可以分析上面定义的f函数。 请注意,它在其输入a中是分段线性的。 换言之,对于任何a,存在某个常量标量k,使得f(a)=k*a,其中k的值取决于输入a,因此可以用d/a验证梯度是否正确。

a_grad == d / a
Tensor(shape=[], dtype=Bool, value= True)

2.5.5. 小结

  • 深度学习框架可以自动计算导数:我们首先将梯度附加到想要对其计算偏导数的变量上,然后记录目标值的计算,执行它的反向传播函数,并访问得到的梯度。

2.5.6. 练习

  1. 为什么计算二阶导数比一阶导数的开销要更大?

  2. 在运行反向传播函数之后,立即再次运行它,看看会发生什么。

  3. 在控制流的例子中,我们计算d关于a的导数,如果将变量a更改为随机向量或矩阵,会发生什么?

  4. 重新设计一个求控制流梯度的例子,运行并分析结果。

  5. 使\(f(x)=\sin(x)\),绘制\(f(x)\)\(\frac{df(x)}{dx}\)的图像,其中后者不使用\(f'(x)=\cos(x)\)