如何使用torch.autograd生成张量值函数的雅可比矩阵？

Question

计算一个函数 f : R^d -> R^d 的雅可比矩阵其实并不难：

def jacobian(y, x):
    k, d = x.shape
    jacobian = list()
    
    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

在上面的代码中，jacobian 是用 y = f(x) 来调用的。不过，现在我有一个函数 g = g(t, x)，其中 t 是一个形状为 k 的 torch.tensor，而 x 是一个形状为 (k, d1, d2, d3) 的 torch.tensor。函数 g 的结果也是一个形状为 (k, d1, d2, d3) 的 torch.tensor。

我尝试使用我已经存在的 jacobian 函数。我做的事情是：

y = g(t, x)
x = x.flatten(1)
y = y.flatten(1)
jacobian(y, x)

问题是，dy_dx 一直是 None。我唯一能想到的解释是，可能在调用 flatten(1) 之后，依赖关系图被破坏了。

那么，我该怎么办呢？我想说明的是，我实际上想计算的是散度，也就是雅可比矩阵的迹。如果有更高效的解决方案来处理这个特定情况，我会对此感兴趣。

Answer 1

你说得对，你传入的其实是 x.flatten(1)，而 y 是从 x 计算出来的，跟 x.flatten(1) 没关系。你可以用下面这种方式避免扁平化处理：

def jacobian(y, x):
    jacobian = []
    for i in range(x.numel()//len(x)):
        v = torch.zeros_like(y.flatten(1))
        v[:, i] = 1.
        dy_dx, *_ = torch.autograd.grad(y, x, v.view_as(x), 
                      retain_graph=True, create_graph=True, allow_unused=True)
        jacobian.append(dy_dx)
    jacobian = torch.stack(jacobian, dim=1)
    return jacobian

这样在调用函数后，你可以把 d1、d2 和 d3 一起扁平化。这里有一个简单的例子：

x = torch.rand(1,3,2,2, requires_grad=True)
J = jacobian(x**2, x)

> J.flatten(-3)
tensor([[[1.2363, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
         [0.0000, 0.2386, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
         [0.0000, 0.0000, 1.0451, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
         [0.0000, 0.0000, 0.0000, 1.4160, 0.0000, 0.0000, 0.0000, 0.0000],
         [0.0000, 0.0000, 0.0000, 0.0000, 0.4090, 0.0000, 0.0000, 0.0000],
         [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 1.7642, 0.0000, 0.0000],
         [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.3041, 0.0000],
         [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 1.1995]]],
       grad_fn=<ViewBackward0>)

接下来，我想你可以用 torch.trace 来计算你想要的迹。

> [j.trace() for j in J.flatten(-3)]
[tensor(7.6128, grad_fn=<TraceBackward0>)]

记得你也可以使用内置的 jacobian 函数：

J = torch.autograd.functional.jacobian(lambda x: g(t, x), x)

不过，你需要对结果进行重新调整形状：

k = x.numel()//len(x)

> [j.trace() for j in J.view(len(x), k, k)]
[tensor(7.6128)]