利用自动微分库计算任意张量的偏导数

Question

(注:这不是一个反向传播问题。)我试图用PyTorch张量来代替Numpy数组来解决GPU上的非线性PDE问题。我要计算任意张量的偏导数，类似于中心有限差分numpy.gradient函数的作用。我还有其他解决这个问题的方法，但是由于我已经在使用PyTorch，我想知道是否可以使用autograd模块(或者，一般情况下，任何其他的自动分化模块)来执行这个操作。

我已经创建了一个与张量兼容的numpy.gradient函数版本，运行速度要快得多。但也许有一种更优雅的方式来做到这一点。我找不到解决这个问题的任何其他来源，也无法证明这是可能的或不可能的；也许这反映了我对自微分算法的无知。

Answer 1

我自己也有过同样的问题:当数值求解PDE时，我们需要一直访问空间梯度( numpy.gradients函数可以给我们)--是否可以使用自动微分来计算梯度，而不是使用有限差分或某种形式的梯度？

“我想知道是否可以使用autograd模块(或者，一般情况下，任何其他的自动区分模块)来执行此操作。”。

答案不是：一旦你在空间或时间中离散你的问题，那么时间和空间就变成一个网格结构的离散变量，而不是显式变量，你把这些变量输入到某个函数中来计算这个问题的解。

例如，如果我想计算某一流体流动u(x，t)的速度场，我会在空间和时间上进行离散，并有u[:,:]，其中指数表示空间和时间中的位置。

自动微分可以计算函数u(x，t)的导数。所以为什么不能在这里计算空间或时间导数呢？因为你已经谨慎地处理了你的问题。这意味着，对于任意x，您没有一个函数，而是在某些网格点上有一个u的函数。对于网格点的间距，您不能自动区分。

据我所知，您编写的张量兼容函数可能是您的最佳选择。您可以在PyTorch论坛、这里和这里中看到类似的问题。或者你可以做这样的事

dx = x[:,:,1:]-x[:,:,:-1]

如果你不担心终点的话。

Answer 2

在某些约束条件下，可以使用PyTorch计算张量相对于另一个张量的梯度。如果您注意保持在张量框架内以确保创建一个计算图，那么通过反复调用输出张量的每个元素并对自变量的梯度成员进行零点化，您可以迭代地查询每个条目的梯度。这种方法允许您逐步构建向量值函数的梯度。

不幸的是，这种方法需要多次调用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)

x0, x1, x2 = (1, 2, 3)的Jacobian可以按以下方式计算

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的一个有用的特性是计算向量雅可比积的能力.前面的示例需要大量链规则(即.a)的重新应用。)通过backward方法直接计算雅可比。但是PyTorch允许用任意向量计算Jacobian的矩阵/向量积，这比实际构造Jacobian要有效得多。这可能更符合您正在寻找的内容，因为您可以在函数的不同值下通过finagle计算多个梯度，这类似于我认为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]