PyTorch 中的雅可比矩阵

介绍：

Jacobian 是一个非常强大的运算符，用于计算给定函数关于其组成潜在变量的偏导数。为了复习的目的，给定函数的雅可比行列式 $f : \R^{n} \rightarrow \R^{m}$ 关于向量 $\mathbf{x} = \{x_1, ..., x_n\} \in \R^{n}$ 被定义为

$\mathbf{J}_f(\mathbf{x}) = \begin{bmatrix} \frac{\partial f}{\partial x_1} & \frac{\partial f}{\partial x_2} & \ldots & \frac{\partial f}{\partial x_n}\end{bmatrix}= \begin{bmatrix} \frac{\partial f_1}{\partial x_1} & \ldots & \frac{\partial f_1}{\partial x_n} \\[1ex] \vdots & \ddots & \vdots \\ \frac{\partial f_m}{\partial x_1} & \ldots & \frac{\partial f_m}{\partial x_n} \end{bmatrix}$

例子：

假设我们有一个向量 $\mathbf{x} = \begin{bmatrix} x_{1} \\[1ex]x_{2} \\[1ex]x_{3} \end{bmatrix}$ 和一个函数 $f(\mathbf{x}) = f(x_1,x_2,x_3)= \begin{bmatrix}f_{1} \\[1ex]f_{2}\\[1ex]f_{3}\end{bmatrix} = \begin{bmatrix}x_1+x_2 \\[1ex]x_1 \times x_3 \\[1ex]x_2^{3}\end{bmatrix}$ .计算雅可比行列式 $f$ 关于 $\mathbf{x}$ ，我们可以用上面的公式得到

$\mathbf{J}_f(\mathbf{x}) = \begin{bmatrix} \frac{\partial f}{\partial x_1} & \frac{\partial f}{\partial x_2} & \frac{\partial f}{\partial x_3}\end{bmatrix}= \begin{bmatrix} \frac{\partial f_1}{\partial x_1} & \frac{\partial f_1}{\partial x_2} & \frac{\partial f_1}{\partial x_3} \\[1ex] \frac{\partial f_2}{\partial x_1} & \frac{\partial f_2}{\partial x_2} & \frac{\partial f_2}{\partial x_3} \\[1ex] \frac{\partial f_3}{\partial x_1} & \frac{\partial f_3}{\partial x_2} & \frac{\partial f_3}{\partial x_3} \end{bmatrix}=\begin{bmatrix} \frac{\partial (x_1 + x_2)}{\partial x_1} & \frac{\partial (x_1 + x_2)}{\partial x_2} & \frac{\partial (x_1 + x_2)}{\partial x_3} \\[1ex] \frac{\partial (x_1 \times x_3)}{\partial x_1} & \frac{\partial (x_1 \times x_3)}{\partial x_2} & \frac{\partial (x_1 \times x_3)}{\partial x_3} \\[1ex] \frac{\partial x_2^3}{\partial x_1} & \frac{\partial x_2^3}{\partial x_2} & \frac{\partial x_2^3}{\partial x_3} \end{bmatrix} = \begin{bmatrix}1 & 1 & 0 \\[1ex] x_3 & 0 & x_1 \\[1ex] 0 & 3\times x_2^2&0\end{bmatrix}$

为了实现与上述相同的功能，我们可以使用 Pytorch 的torch.autograd.functional实用程序中的jacobian()函数来计算给定函数的某些输入的雅可比矩阵。

Syntax: torch.autograd.functional.jacobian(func, inputs, create_graph=False, strict=False, vectorize=False)

Parameters:

func: A Python function which takes input and outputs a Pytorch Tensor(or a tuple of Tensors).
inputs: The inputs are passed as parameters to the ‘func’ method. The input can be a single Pytorch Tensor(or a tuple of Tensors)
create_graph: If True, the autograd engine creates a backpropable graph for doing further operations on gradients. Defaults to False.
strict: If True, an error will be raised when the engine detects that there exists an input such that all the outputs are independent of it. If False, returns zero gradients for such inputs. Defaults to False.
vectorize: Still in it’s experimental phase if True, the function uses the vmap prototype feature to compute the gradients by only one call of the autograd engine instead of one call per each row of the matrix. Defaults to False.

编程需要懂一点英语

安装：

对于本文，您只需要 Torch 实用程序，可以使用以下命令通过 pip 包管理器下载：

pip install torch

函数的使用示例：

我们将使用相同的函数和向量来简化流程，如上例所述。由于张量是 Pytorch 包的基本构建块，我们将使用它们来表示输入向量和给定函数。本文假设您对 Pytorch 张量有基本的了解，可以通过阅读 Pytorch 文章快速复习。

理论验证：

假设我们有一个向量 $\mathbf{x} = \begin{bmatrix} x_{1} \\[1ex]x_{2} \\[1ex]x_{3} \end{bmatrix} = \begin{bmatrix} 3 \\[1ex]4 \\[1ex]5 \end{bmatrix}$ 作为给定的输入。通过插入值 $\mathbf{x}$ 代入上面的推导方程，我们将得到 $J_f({\mathbf{x}}) =\begin{bmatrix}1 & 1 & 0 \\[1ex] x_3 & 0 & x_1 \\[1ex] 0 & 3\times x_2^2&0\end{bmatrix}=\begin{bmatrix}1 & 1 & 0 \\[1ex] 5 & 0 & 3 \\[1ex] 0 & 3\times 4^2&0\end{bmatrix}= \begin{bmatrix}1 & 1 & 0 \\[1ex] 5 & 0 & 3 \\[1ex] 0 & 48 & 0\end{bmatrix}$

代码：使用 Pytorch 显示雅可比矩阵工作的Python实现

Python

from torch.autograd.functional import jacobian
from torch import tensor
 
#Defining the main function
def f(x1,x2,x3):
    return (x1 + x2, x3*x1, x2**3)
 
#Defining input tensors
x1 = tensor(3.0)
x2 = tensor(4.0)
x3 = tensor(5.0)
 
#Printing the Jacobian
print(jacobian(f,(x1,x2,x3)))

输出：

((tensor(1.), tensor(1.), tensor(0.)), 
(tensor(5.), tensor(0.), tensor(3.)), 
(tensor(0.), tensor(48.), tensor(0.)))

输出与我们的理论验证完全相似！使用类似的方法，我们可以使用 Pytorch API 计算任何给定函数的雅可比矩阵。

参考：

https://pytorch.org/docs/stable/autograd.html#torch.autograd.functional.jacobian
https://pytorch.org/docs/stable/tensors.html