使用函数变换计算雅可比矩阵和黑塞矩阵
雅可比矩阵
在介绍MindSpore提供的计算雅可比矩阵的方法之前，首先对雅可比矩阵进行介绍。
我们首先定义一个映射\(\textbf{f}\):
我们在这里使用符号\(\longmapsto\)表示集合元素之间的映射，并加粗所有代表向量的符号，因此有：
其中\(\textbf{x} = (x_{1}, x_{2},\dots, x_{n})\)，\(\textbf{f(x)} = (f_{1}(\textbf{x}), f_{2}(\textbf{x}),\dots,f_{m}(\textbf{x}))\)。
我们将所有从\(R^{n}\)到\(R^{m}\)的映射构成的集合记为\(F_{n}^{m}\)。 在这里，我们将从一个函数（映射）集到另一个函数（映射）集的映射称为一个操作。易得，梯度操作\(\nabla\)是在函数集\(F_{n}^{1}\)上同时进行n次偏导数的操作，将其定义为一个从函数集\(F_{n}^{1}\)到\(F_{n}^{n}\)的映射：
广义的梯度操作\(\partial\)，被定义为在函数集\(F_{n}^{m}\)上同时进行n次偏导数的操作，
雅可比矩阵就是将操作\(\partial\)作用于\(\textbf{f}\)后得到的结果，即
得到雅可比矩阵，
雅可比矩阵的应用：在自动微分的前向模式中，每一次前向传播，可以求出雅可比矩阵的一列。在自动微分的反向模式中，每一次反向传播，我们可以计算出雅可比矩阵的一行。
计算雅可比矩阵
使用标准的自动微分系统很难高效地计算雅可比矩阵，但MindSpore却提供了能够高效计算雅可比的方法，下面对这些方法进行介绍。
首先，我们定义一个函数forecast，该函数是一个简单的线性函数，并带有一个非线性激活函数。
[1]:
import time
import mindspore
from mindspore import ops
from mindspore import jacrev, jacfwd, vmap, vjp, jvp, grad
import numpy as np
mindspore.set_seed(1)
def forecast(weight, bias, x):
return ops.dense(x, weight, bias).tanh()
接下来，我们构造一些数据：一个权重张量weight，一个偏差张量bias，还有一个输入向量x。
[2]:
D = 16
weight = ops.randn(D, D)
bias = ops.randn(D)
x = ops.randn(D)
函数forecast对输入向量x做如下映射变换，\(R^{D}\overset{}{\rightarrow}R^{D}\)。 MindSpore在自动微分的过程中，会计算向量-雅可比积。为了计算映射\(R^{D}\overset{}{\rightarrow}R^{D}\)的完整雅可比矩阵，我们每次使用不同的单位向量逐行来计算它。
[3]:
def partial_forecast(x):
return ops.dense(x, weight, bias).tanh()
_, vjp_fn = vjp(partial_forecast, x)
def compute_jac_matrix(unit_vectors):
jacobian_rows = [vjp_fn(vec)[0] for vec in unit_vectors]
return ops.stack(jacobian_rows)
unit_vectors = ops.eye(D)
jacobian = compute_jac_matrix(unit_vectors)
print(jacobian.shape)
print(jacobian[0])
(16, 16)
[-3.2045446e-05 -1.3530695e-05 1.8671712e-05 -9.6547810e-05
5.9755850e-05 -5.1343523e-05 1.3528993e-05 -4.6988782e-05
-4.5517798e-05 -6.1188715e-05 -1.6264191e-04 5.5033437e-05
-4.3497541e-05 2.2357668e-05 -1.3188722e-04 -3.0677278e-05]
在compute_jac_matrix中，使用for循环逐行计算的方式计算雅可比矩阵，计算效率并不高。MindSpore提供jacrev来计算雅可比矩阵，jacrev的实现利用了vmap，vmap可以消除compute_jac_matrix中的for循环并向量化整个计算过程。jacrev的参数grad_position指定计算输出相对于哪个参数的雅可比矩阵。
[4]:
from mindspore import jacrev
jacrev_jacobian = jacrev(forecast, grad_position=2)(weight, bias, x)
assert np.allclose(jacrev_jacobian.asnumpy(), jacobian.asnumpy())
接下来对compute_jac_matrix和jacrev的性能进行对比。通常情况下，jacrev的性能更好，因为使用vmap进行向量化运算，会更充分地利用硬件，同时计算多组数据，降低计算开销，以获得更好的性能。
让我们编写一个函数，在微秒量级上评估两种方法的性能。
[5]:
def perf_compution(func, run_times, *args, **kwargs):
start_time = time.perf_counter()
for _ in range(run_times):
func(*args, **kwargs)
end_time = time.perf_counter()
cost_time = (end_time - start_time) * 1000000
return cost_time
run_times = 500
xp = x.copy()
compute_jac_matrix_cost_time = perf_compution(compute_jac_matrix, run_times, xp)
jac_fn = jacrev(forecast, grad_position=2)
jacrev_cost_time = perf_compution(jac_fn, run_times, weight, bias, x)
print(f"compute_jac_matrix run {run_times} times, cost time {compute_jac_matrix_cost_time} microseconds.")
print(f"jacrev run {run_times} times, cost time {jacrev_cost_time} microseconds.")
compute_jac_matrix run 500 times, cost time 12942823.04868102 microseconds.
jacrev run 500 times, cost time 909309.7001314163 microseconds.
分别运行compute_jac_matrix和jacrev500次，统计它们消耗的时间。
下面计算，相较于使用compute_jac_matrix，使用jacrev计算雅可比矩阵，性能提升的百分比。
[6]:
def perf_cmp(first, first_descriptor, second, second_descriptor):
faster = second
slower = first
gain = (slower - faster) / slower
if gain < 0:
gain *= -1
final_gain = gain*100
print(f" Performance delta: {final_gain:.4f} percent improvement with {second_descriptor}. ")
perf_cmp(compute_jac_matrix_cost_time, "for loop", jacrev_cost_time, "jacrev")
Performance delta: 92.9744 percent improvement with jacrev.
此外，也可以通过指定jacrev的参数grad_position来计算输出相对于模型参数weight和bias的雅可比矩阵。
[7]:
jacrev_weight, jacrev_bias = jacrev(forecast, grad_position=(0, 1))(weight, bias, x)
print(jacrev_weight.shape)
print(jacrev_bias.shape)
(16, 16, 16)
(16, 16)
反向模式计算雅可比矩阵 vs 前向模式计算雅可比矩阵
MindSpore提供了两个API来计算雅可比矩阵：分别是jacrev和jacfwd。
jacrev：使用反向模式自动微分。jacfwd：使用前向模式自动微分。
jacfwd和jacrev可以相互替换，但是它们在不同的场景下，性能表现不同。
一般来说，如果需要计算函数\(R^{n}\overset{}{\rightarrow}R^{m}\)的雅可比矩阵，当该函数的输出向量的规模大于输入向量的规模时（即，m > n），jacfwd在性能方面表现得更好，否则，jacrev在性能方面表现得更好。
下面对这个结论做一个非严谨的论证，在前向模式自动微分（计算雅可比-向量积）的过程中，是逐列计算雅可比矩阵的，在反向模式自动微分（计算向量-雅可比积）的过程中，是逐行计算雅可比矩阵的。假设待计算的雅可比矩阵的规模是m行，n列，如果m > n，我们推荐使用逐列计算雅可比矩阵的jacfwd，反之，如果m < n，我们推荐使用逐行计算雅可比矩阵的jacrev。
黑塞矩阵
在介绍MindSpore提供的计算黑塞矩阵的方法之前，首先对黑塞矩阵进行介绍。
黑塞矩阵可以由梯度操作\(\nabla\)和广度梯度操作\(\partial\)的复合得到，即
将该复合操作用于f，得到，
可以得到黑塞矩阵，
易见，黑塞矩阵是一个实对称矩阵。
黑塞矩阵的应用：利用黑塞矩阵，我们可以探索神经网络在某点处的曲率，为训练是否收敛提供数值依据。
计算黑塞矩阵
在MindSpore中，我们可以通过jacfwd和jacrev的任意组合来计算黑塞矩阵。
[8]:
Din = 32
Dout = 16
weight = ops.randn(Dout, Din)
bias = ops.randn(Dout)
x = ops.randn(Din)
hess1 = jacfwd(jacfwd(forecast, grad_position=2), grad_position=2)(weight, bias, x)
hess2 = jacfwd(jacrev(forecast, grad_position=2), grad_position=2)(weight, bias, x)
hess3 = jacrev(jacfwd(forecast, grad_position=2), grad_position=2)(weight, bias, x)
hess4 = jacrev(jacrev(forecast, grad_position=2), grad_position=2)(weight, bias, x)
np.allclose(hess1.asnumpy(), hess2.asnumpy())
np.allclose(hess2.asnumpy(), hess3.asnumpy())
np.allclose(hess3.asnumpy(), hess4.asnumpy())
[8]:
True
计算批量雅可比矩阵和批量黑塞矩阵
在上面给出的示例中，我们都是计算单一的输出向量关于单一的输入向量的雅可比矩阵。在某些情况下，你可能想计算一批量的输出向量关于一批量的输入向量的雅可比矩阵，或者换句话说，给出一批量的输入向量，其shape为(b, n)，函数的映射关系是\(R^{n}\overset{}{\rightarrow}R^{m}\)，我们期望得到一批量的雅可比矩阵，其shape为(b, m, n)。
我们可以使用vmap计算批量雅可比矩阵。
[9]:
batch_size = 64
Din = 31
Dout = 33
weight = ops.randn(Dout, Din)
bias = ops.randn(Dout)
x = ops.randn(batch_size, Din)
compute_batch_jacobian = vmap(jacrev(forecast, grad_position=2), in_axes=(None, None, 0))
batch_jacobian = compute_batch_jacobian(weight, bias, x)
print(batch_jacobian.shape)
(64, 33, 31)
计算批量黑塞矩阵的方式与计算批量雅可比矩阵的方法类似，也可以使用vmap来计算批量黑塞矩阵，
[10]:
hessian = jacrev(jacrev(forecast, grad_position=2), grad_position=2)
compute_batch_hessian = vmap(hessian, in_axes=(None, None, 0))
batch_hessian = compute_batch_hessian(weight, bias, x)
print(batch_hessian.shape)
(64, 33, 31, 31)
计算黑塞-向量积
计算黑塞-向量积（Hessian-vector product, hvp）的最直接的方法计算一个完整的黑塞矩阵，并将其与向量进行点积运算。但MindSpore提供了更好的方法，使得不需要计算一个完整的黑塞矩阵，便可以计算黑塞-向量积。下面我们介绍计算黑塞-向量积的两种方法。
将反向模式自动微分与反向模式自动微分组合。
将反向模式自动微分与前向模式自动微分组合。
下面先介绍，在MindSpore中，如何使用反向模式自动微分与前向模式自动微分组合的方式计算黑塞-向量积，
[11]:
def hvp_revfwd(f, inputs, vector):
return jvp(grad(f), inputs, vector)[1]
def f(x):
return x.sin().sum()
inputs = ops.randn(128)
vector = ops.randn(128)
result_hvp_revfwd = hvp_revfwd(f, inputs, vector)
print(result_hvp_revfwd.shape)
(128,)
如果前向自动微分不能满足要求，我们可以使用反向模式自动微分与反向模式自动微分组合的方式来计算黑塞-向量积，
[12]:
def hvp_revrev(f, inputs, vector):
_, vjp_fn = vjp(grad(f), *inputs)
return vjp_fn(*vector)
result_hvp_revrev = hvp_revrev(f, (inputs,), (vector,))
print(result_hvp_revrev[0].shape)
(128,)
使用上面两种方法计算黑塞-向量积得到的结果是一样的。
[13]:
assert np.allclose(result_hvp_revfwd.asnumpy(), result_hvp_revrev[0].asnumpy())