相关疑难解决方法(0)

PyTorch 最高效的 Jacobian/Hessian 计算

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

# 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()
Run Code Online (Sandbox Code Playgroud)
# 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
Run Code Online (Sandbox Code Playgroud)
# 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
Run Code Online (Sandbox Code Playgroud)

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

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

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

python linear-algebra calculus deep-learning pytorch

5
推荐指数
1
解决办法
1385
查看次数