mindspore.jacrev
- mindspore.jacrev(fn, grad_position=0, has_aux=False)[源代码]
通过反向模式计算给定网络的雅可比矩阵,对应 反向模式自动微分。当网络输出数量远小于输入数量时,使用反向模式求雅可比矩阵比前向模式性能更好。
- 参数:
fn (Union[Cell, Function]) - 待求导的函数或网络。以Tensor为入参,返回Tensor或Tensor数组。
grad_position (Union[int, tuple[int]], 可选) - 指定求导输入位置的索引。若为int类型,表示对单个输入求导;若为tuple类型,表示对tuple内索引的位置求导,其中索引从0开始。默认值:
0
。has_aux (bool, 可选) - 若 has_aux 为
True
,只有 fn 的第一个输出参与 fn 的求导,其他输出将直接返回。此时, fn 的输出数量必须超过一个。默认值:False
。
- 返回:
Function,用于计算给定函数的雅可比矩阵。例如 out1, out2 = fn(*args) ,若 has_aux 为
True
,梯度函数将返回 (Jacobian, out2) 形式的结果,其中 out2 不参与求导,若为False
,将直接返回 Jacobian 。- 异常:
TypeError - grad_position 或 has_aux 类型不符合要求。
- 支持平台:
Ascend
GPU
CPU
样例:
>>> import numpy as np >>> import mindspore.nn as nn >>> from mindspore import jacrev >>> from mindspore import Tensor >>> class MultipleInputsMultipleOutputsNet(nn.Cell): ... def construct(self, x, y, z): ... return x ** 2 + y ** 2 + z ** 2, x * y * z >>> x = Tensor(np.array([[1, 2], [3, 4]]).astype(np.float32)) >>> y = Tensor(np.array([[1, 2], [3, 4]]).astype(np.float32)) >>> z = Tensor(np.array([[1, 1], [1, 1]]).astype(np.float32)) >>> net = MultipleInputsMultipleOutputsNet() >>> jac, aux = jacrev(net, grad_position=0, has_aux=True)(x, y, z) >>> print(jac) [[[[ 2. 0.] [ 0. 0.]] [[ 0. 4.] [ 0. 0.]]] [[[ 0. 0.] [ 6. 0.]] [[ 0. 0.] [ 0. 8.]]]] >>> print(aux) [[ 1. 4.] [ 9. 16.]]