昇思MindSpore进阶教程--雅可比矩阵

大家好，我是刘明，明志科技创始人，华为昇思MindSpore布道师。
技术上主攻前端开发、鸿蒙开发和AI算法研究。
努力为大家带来持续的技术分享，如果你也喜欢我的文章，就点个关注吧

雅可比矩阵

雅可比矩阵的应用：在自动微分的前向模式中，每一次前向传播，可以求出雅可比矩阵的一列。在自动微分的反向模式中，每一次反向传播，我们可以计算出雅可比矩阵的一行。

计算雅可比矩阵

使用标准的自动微分系统很难高效地计算雅可比矩阵，但MindSpore却提供了能够高效计算雅可比的方法，下面对这些方法进行介绍。

首先，我们定义一个函数forecast，该函数是一个简单的线性函数，并带有一个非线性激活函数。

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。

D = 16
weight = ops.randn(D, D)
bias = ops.randn(D)
x = ops.randn(D)

函数forecast对输入向量x做如下映射变换，$R^D\to R^D$。 MindSpore在自动微分的过程中，会计算向量-雅可比积。为了计算映射$R^D\to R^D$的完整雅可比矩阵，我们每次使用不同的单位向量逐行来计算它。

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])

在compute_jac_matrix中，使用for循环逐行计算的方式计算雅可比矩阵，计算效率并不高。MindSpore提供jacrev来计算雅可比矩阵，jacrev的实现利用了vmap，vmap可以消除compute_jac_matrix中的for循环并向量化整个计算过程。jacrev的参数grad_position指定计算输出相对于哪个参数的雅可比矩阵。

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进行向量化运算，会更充分地利用硬件，同时计算多组数据，降低计算开销，以获得更好的性能。

让我们编写一个函数，在微秒量级上评估两种方法的性能。

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和jacrev500次，统计它们消耗的时间。

下面计算，相较于使用compute_jac_matrix，使用jacrev计算雅可比矩阵，性能提升的百分比。

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")

此外，也可以通过指定jacrev的参数grad_position来计算输出相对于模型参数weight和bias的雅可比矩阵。

jacrev_weight, jacrev_bias = jacrev(forecast, grad_position=(0, 1))(weight, bias, x)
print(jacrev_weight.shape)
print(jacrev_bias.shape)

反向模式计算雅可比矩阵 vs 前向模式计算雅可比矩阵

MindSpore提供了两个API来计算雅可比矩阵：分别是jacrev和jacfwd。

• jacrev：使用反向模式自动微分。

• jacfwd：使用前向模式自动微分。

jacfwd和jacrev可以相互替换，但是它们在不同的场景下，性能表现不同。

一般来说，如果需要计算函数
的雅可比矩阵，当该函数的输出向量的规模大于输入向量的规模时（即，m > n），jacfwd在性能方面表现得更好，否则，jacrev在性能方面表现得更好。

下面对这个结论做一个非严谨的论证，在前向模式自动微分（计算雅可比-向量积）的过程中，是逐列计算雅可比矩阵的，在反向模式自动微分（计算向量-雅可比积）的过程中，是逐行计算雅可比矩阵的。假设待计算的雅可比矩阵的规模是m行，n列，如果m > n，我们推荐使用逐列计算雅可比矩阵的jacfwd，反之，如果m < n，我们推荐使用逐行计算雅可比矩阵的jacrev。