张量函数雅可比行列式计算详解：PyTorch.autograd 助力

python

2024-03-17 05:01:54

掌握张量函数雅可比行列式计算，PyTorch.autograd 助力

简介

在机器学习和深度学习领域，计算张量函数的雅可比行列式是一项至关重要的任务。雅可比行列式揭示了函数导数的矩阵形式，对于理解函数行为和解决优化问题不可或缺。

借助 PyTorch.autograd

PyTorch.autograd 是一个自动微分模块，允许我们轻松计算张量的导数。它可以通过 torch.autograd.grad 函数针对给定的输入张量自动求导数。

计算张量值函数的雅可比行列式

函数 f(x)

考虑一个从实数域 R^d 映射到 R^d 的函数 f(x)。计算其雅可比行列式非常简单：

def jacobian(y, x):
    k, d = x.shape
    jacobian = []
    for i in range(d):
        v = torch.zeros_like(y)
        v[:, i] = 1.
        dy_dx = torch.autograd.grad(y, x, grad_outputs=v, retain_graph=True, create_graph=True, allow_unused=True)[0]  # shape [k, d]
        jacobian.append(dy_dx)
    jacobian = torch.stack(jacobian, dim=1).requires_grad_()
    return jacobian

函数 g(t, x)

对于一个函数 g(t, x)，其中 t 是一个形状为 k 的张量，x 是一个形状为 (k, d1, d2, d3) 的张量，雅可比行列式的计算可以通过修改 jacobian 函数来实现：

# 需要先将输入张量扁平化
x = x.flatten(1)
y = y.flatten(1)
# 使用 jacobian 计算
jacobian(y, x)

替代方法：计算散度

如果你只对雅可比行列式的迹（散度）感兴趣，那么可以采用更有效的方法：

def divergence(y, x):
    k, d1, d2, d3 = x.shape
    y = y.flatten(1)
    grad = torch.autograd.grad(y, x, create_graph=True, allow_unused=True)
    div = 0.
    for i in range(d1 * d2 * d3):
        div += grad[0][:, i, i].sum()
    return div