Question

（注意：这不是有关反向传播的问题。）我正在尝试在GPU上使用PyTorch张量代替Numpy数组解决非线性PDE问题。我想计算任意张量的偏导数，类似于中心有限差分numpy.gradient函数的作用。我还有其他方法可以解决此问题，但是由于我已经在使用PyTorch，所以我想知道是否可以使用autograd模块（或者通常是任何其他自动分化模块）来执行此操作。

我创建了numpy.gradient函数的张量兼容版本-运行速度快得多。但是，也许有更优雅的方法可以做到这一点。我找不到其他任何资料可以解决这个问题，无论是表明这是可能的还是不可能的。也许这反映了我对自动分化算法的无知。

Answer 1

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

“我想知道是否可以使用autograd模块（或者通常是任何其他自动分化模块）来执行此操作。”

答案是否定的：一旦您在空间或时间上离散问题，时间和空间就会变成具有网格状结构的离散变量，而不是您要输入的显式变量一些函数来计算PDE的解。

例如，如果我想计算某些流体流u（x，t）的速度场，我将在空间和时间上离散化，并且我将拥有u[:,:]，其中的索引表示空间中的位置，时间。

自动微分可以计算函数u（x，t）的导数。那么为什么它不能在这里计算空间或时间导数呢？因为您离散化了您的问题。这意味着您没有任意x的u函数，而是在某些网格点具有u函数。您无法根据网格点的间距自动区分。

据我所知，您编写的与张量兼容的函数可能是最好的选择。您可以在PyTorch论坛here和here中看到类似的问题。或者您可以做类似

的操作

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

如果您不担心端点。

Answer 2

在某些约束下，您可以使用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)

x0, 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的一个有用功能是能够计算vector-Jacobian乘积。前面的示例需要通过backward方法重新应用链式规则（也称为反向传播），以直接计算雅可比行列式。但是PyTorch允许您使用任意向量来计算Jacobian的矩阵/向量乘积，这比实际构建Jacobian的效率要高得多。这可能更符合您的需求，因为您可以将其锁定以在函数的各个值上计算多个梯度，类似于我相信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]

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

2 个答案: