gpt4 book ai didi

python - PyTorch 最高效的 Jacobian/Hessian 计算

转载 作者:行者123 更新时间:2023-12-02 09:24:00 31 4
gpt4 key购买 nike

我正在寻找通过 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()
# 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 中计算雅可比行列式。

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

最后,有谁知道在 TensorFlow 中是否可以更轻松或更高效地完成类似的事情?

最佳答案

functorch 可以进一步加快计算速度。例如,此代码来自 functorch docs用于批量雅可比计算(Hessian 也适用):

batch_size = 64
Din = 31
Dout = 33

weight = torch.randn(Dout, Din)
print(f"weight shape = {weight.shape}")
bias = torch.randn(Dout)

def predict(weight, bias, x):
return F.linear(x, weight, bias).tanh()

x = torch.randn(batch_size, Din)
compute_batch_jacobian = vmap(jacrev(predict, argnums=2), in_dims=(None, None, 0))
batch_jacobian0 = compute_batch_jacobian(weight, bias, x)

关于python - PyTorch 最高效的 Jacobian/Hessian 计算,我们在Stack Overflow上找到一个类似的问题: https://stackoverflow.com/questions/56480578/

31 4 0
Copyright 2021 - 2024 cfsdn All Rights Reserved 蜀ICP备2022000587号
广告合作:1813099741@qq.com 6ren.com