您可以使用 PyTorch 在某些约束下计算一个张量相对于另一个张量的梯度。如果您小心地留在张量框架内以确保创建计算图,那么通过在输出张量的每个元素上重复调用并将自变量的 grad 成员归零,您可以迭代地查询每个条目的梯度。这种方法允许您逐步构建向量值函数的梯度。
backward
不幸的是,这种方法需要多次调用,这在实践中可能很慢并且可能导致非常大的矩阵。
import torch
from copy import deepcopy
def get_gradient(f, x):
""" computes gradient of tensor f with respect to tensor x """
assert x.requires_grad
x_shape = x.shape
f_shape = f.shape
f = f.view(-1)
x_grads = []
for f_val in f:
if x.grad is not None:
x.grad.data.zero_()
f_val.backward(retain_graph=True)
if x.grad is not None:
x_grads.append(deepcopy(x.grad.data))
else:
# in case f isn't a function of x
x_grads.append(torch.zeros(x.shape).to(x))
output_shape = list(f_shape) + list(x_shape)
return torch.cat((x_grads)).view(output_shape)
例如,给定以下函数:
f(x0,x1,x2) = (x0*x1*x2, x1^2, x0+x2)
雅可比 atx0, x1, x2 = (1, 2, 3)
可以计算如下
x = torch.tensor((1.0, 2.0, 3.0))
x.requires_grad_(True) # must be set before further computation
f = torch.stack((x[0]*x[1]*x[2], x[1]**2, x[0]+x[2]))
df_dx = get_gradient(f, x)
print(df_dx)
这导致
tensor([[6., 3., 2.],
[0., 4., 0.],
[1., 0., 1.]])
对于您的情况,如果您可以相对于输入张量定义输出张量,则可以使用这样的函数来计算梯度。
PyTorch 的一个有用特性是能够计算向量-雅可比积。backward
前面的示例需要通过直接计算雅可比行列的方法对链式法则(也称为反向传播)进行大量重新应用。但是 PyTorch 允许您使用任意向量计算雅可比矩阵/向量乘积,这比实际构建雅可比行列式要高效得多。这可能更符合您正在寻找的内容,因为您可以使用它来计算函数不同值的多个梯度,类似于我认为的numpy.gradient
操作方式。
例如,在这里我们计算f(x) = x^2 + sqrt(x)
并计算每个点x = 1, 1.1, ..., 1.8
的导数(即)f'(x) = 2x + 0.5/sqrt(x)
dx = 0.1
x = torch.arange(1, 1.8, dx, requires_grad=True)
f = x**2 + torch.sqrt(x)
f.backward(torch.ones(f.shape))
x_grad = x.grad
print(x_grad)
这导致
tensor([2.5000, 2.6767, 2.8564, 3.0385, 3.2226, 3.4082, 3.5953, 3.7835])
将此与 numpy.gradient 进行比较
dx = 0.1
x_np = np.arange(1, 1.8, dx)
f_np = x_np**2 + np.sqrt(x_np)
x_grad_np = np.gradient(f_np, dx)
print(x_grad_np)
这导致以下近似值
[2.58808848 2.67722558 2.85683288 3.03885421 3.22284723 3.40847554 3.59547805 3.68929417]