PyTorch 最高效的 Jacobian/Hessian 计算 [英] PyTorch most efficient Jacobian/Hessian calculation

问题描述

我正在寻找通过 Pytorch 获取函数的雅可比行列式的最有效方法，并且到目前为止提出了以下解决方案:

I am looking for the most efficient way to get the Jacobian of a function through Pytorch and have so far come up with the following solutions:

# Setup
def func(X):
    return torch.stack((X.pow(2).sum(1),
                        X.pow(3).sum(1),
                        X.pow(4).sum(1)),1)  

X = Variable(torch.ones(1,int(1e5))*2.00094, requires_grad=True).cuda()

# Solution 1:
t = time()
Y = func(X)
J = torch.zeros(3, int(1e5))

for i in range(3):
    J[i] = grad(Y[0][i], X, create_graph=True, retain_graph=True, allow_unused=True)[0]

print(time()-t)
>>> Output: 0.002 s

# Solution 2:
def Jacobian(f,X):
    X_batch = Variable(X.repeat(3,1), requires_grad=True)
    f(X_batch).backward(torch.eye(3).cuda(), retain_graph=True)
    return X_batch.grad

t = time()
J2 = Jacobian(func,X)
print(time()-t)
>>> Output: 0.001 s

由于在第一个解决方案中使用循环与在第二个解决方案中使用循环似乎没有太大区别，我想问一下是否还有一种更快的方法来计算 pytorch 中的雅可比行列式.

Since there seem to be not a big difference between using a loop in the first solution than the second one, I wanted to ask if there might still be be a faster way to calculate a Jacobian in pytorch.

我的另一个问题也是关于计算 Hessian 的最有效方法是什么.

My other question is then also about what might be the most efficient way to calculate the Hessian.

最后，有谁知道在 TensorFlow 中是否可以更轻松或更高效地完成此类操作?

Finally, does anyone know if something like this can be done easier or more efficient in TensorFlow?

推荐答案

最有效的方法可能是使用 PyTorch 自带的内置函数:

The most efficient method is likely to use PyTorch's own inbuilt functions:

torch.autograd.functional.jacobian(func, x)
torch.autograd.functional.hessian(func, x)

这篇关于PyTorch 最高效的 Jacobian/Hessian 计算的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！