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:

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中计算Jacobian.

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?

推荐答案

我遇到了一个类似的问题，该问题通过手动定义Jacobian(手动计算导数)来解决.对于我的问题，这是可行的，但我可以想象并非总是如此.与第二种解决方案相比，计算时间加快了我计算机(cpu)上的某些因素.

I had a similar problem which I solved by defining the Jacobian manually (calculating the derivatives by hand). For my problem this was feasible, but I can imagine that is not always the case. The computation time speeds up some factors on my machine (cpu), compared to the second solution.

# 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

%timeit Jacobian(func,X)
11.7 ms ± 130 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

# Solution 3
def J_func(X):
    return torch.stack(( 
                 2*X,
                 3*X.pow(2),
                 4*X.pow(3)
                  ),1)

%timeit J_func(X)
539 µs ± 24.4 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

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