获取输出相对于输入的梯度

Question

我目前正在尝试使用 Pytorch 实现 ODE 求解器，我的解决方案需要计算每个输出 wtr 与其输入的梯度。

y = model(x)

for i in range(len(y)): #compute output grad wrt input
       y[i].backward(retain_graph=True)
    
ydx=x.grad

我想知道是否有更优雅的方法来计算批处理中每个输出的梯度，因为代码对于高阶 ODE 和 PDE 会变得混乱。我尝试使用：

torch.autograd.backward(x,y,retain_graph=True)

没有太大的成功。

Answer 1

可以直接使用torch.autograd.grad函数获取梯度。一个问题是它要求输出 (y) 是标量。由于您的输出是一个数组，您仍然需要遍历它的值。

调用看起来像这样。

[torch.autograd.grad(outputs=out, inputs=x, retain_graph=True)[0][i] 
    for i, out in enumerate(y)]

这是我的意思的一个例子。让我们考虑具有值 [1,2,3] 的变量 x 和一个仅对其输入求平方的模型。

x = torch.Tensor([1, 2, 3])
x.requires_grad = True

def model(x): return x ** 2

y = model(x)

现在，如果您按照我的描述调用 torch.autograd.grad，您将得到：

[torch.autograd.grad(outputs=out, inputs=x, retain_graph=True)[0][i] 
    for i, out in enumerate(y)]

# [tensor(2.), tensor(4.), tensor(6.)]

这是衍生品列表。 x - [ydx0, ydx1, ydx2]

的值

Answer 2

如果您的 PyTorch 版本实现了 API，请尝试 torch.autograd.functional.jacobian。我正在为 Burger 方程做同样的事情，并在同一主题上发布了这个帖子：PyTorch how to compute second order jacobian?

用 DL 求解 PDE 是目前的热门话题

Getting the output's grad with respect to the input