1. Jacobian matrix

• 计算
  $$\begin{array}{c}f(x)=\left[\begin{array}{c}{f}_1=x_1^2+2x_2\\ \\ f_2=3x_1+4x_2^2\end{array}\right],J=\left[\begin{array}{cc}\frac{\partial f_1}{\partial x_1}& \frac{\partial f_1}{\partial x_2}\\ & \\ \frac{\partial f_2}{\partial x_1}& \frac{\partial f_2}{\partial x_2}\end{array}\right]=\left[\begin{array}{cc}2x_1& 2\\ & \\ 3& 8x_2\end{array}\right]\end{array}$$
• 我们假设$x_1=1,x_2=2$,可得Jacobian matrix 表示如下：
  $$\begin{array}{c}J=\left[\begin{array}{cc}2x_1& 2\\ & \\ 3& 8x_2\end{array}\right]=\left[\begin{array}{cc}2& 2\\ & \\ 3& 16\end{array}\right]\end{array}$$
• 因为fx为向量，所以在pytorch中一般是没有的，我们需要定义一个标量z
  $$\begin{array}{c}z=\left[\begin{array}{cc}1& 1\end{array}\right]\left[\begin{array}{c}{f}_1\\ \\ f_2\end{array}\right]\to \frac{\partial z}{\partial f}=\left[\begin{array}{cc}1& 1\end{array}\right]\end{array}$$
• 根据链式法则：
  $$\begin{array}{c}\frac{\partial z}{\partial f}\cdot \frac{\partial f}{\partial x}=\frac{\partial z}{\partial x}\end{array}$$
• 假设我们代入可得：
  $$\begin{array}{c}\frac{\partial z}{\partial f}\cdot \frac{\partial f}{\partial x}=\left[\begin{array}{cc}1& 1\end{array}\right]\left[\begin{array}{cc}2& 2\\ & \\ 3& 16\end{array}\right]=\left[\begin{array}{cc}5& 18\end{array}\right]\end{array}$$
• 根据pytorch中的自动微分，我们可以得到
  $$\begin{array}{c}\frac{\partial z}{\partial x}=\left[\begin{array}{cc}5& 18\end{array}\right]\end{array}$$
• 小结1：我们用自动微分，是无法像直接用jacobian函数样求解jacobian矩阵，而是因为引入[1,1]向量后，只能求得和值[5,18]
• 小结2：想用反向传播求得jacobian矩阵，需要将fx中的每一项标量拆开求得，后再叠加在一起

2. python 代码

#!/usr/bin/env python
# -*- coding:utf-8 -*-
# @FileName  :JacobianTest.py
# @Time      :2024/11/25 19:28
# @Author    :Jason Zhang
import torch
from torch import nn
from torch.autograd.functional import jacobian


def funx(x):
    return torch.stack([
        x[0] ** 2 + 2 * x[1], 3 * x[0] + 4 * x[1] ** 2
    ])


if __name__ == "__main__":
    run_code = 0
    in_x = torch.tensor([1.0, 2.0], dtype=torch.float, requires_grad=True)
    y = funx(in_x)
    y.backward(torch.ones_like(y))
    my_jacobian = jacobian(funx, in_x)
    print(f"my_jacobian=\n{my_jacobian}")
    x_grad = in_x.grad
    in_x_ones = torch.ones_like(in_x)
    jacobian_x_grad = in_x_ones @ my_jacobian
    print(f"x_grad={x_grad}")
    print(f"jacobian_x_grad={jacobian_x_grad}")

my_jacobian=
tensor([[ 2.,  2.],
        [ 3., 16.]])
x_grad=tensor([ 5., 18.])
jacobian_x_grad=tensor([ 5., 18.])