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

Question

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

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

Answer 1

您可以使用 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 的一个有用功能是能够计算向量-雅可比积。前面的示例需要通过 backward 方法多次重新应用链式法则（a.k.a. 反向传播）来直接计算雅可比行列式。但是 PyTorch 允许您使用任意向量计算雅可比行列式的 matrix/vector 乘积，这比实际构建雅可比行列式要高效得多。这可能更符合你正在寻找的东西，因为你可以欺骗它来计算函数的不同值的多个梯度，类似于我相信 numpy.gradient 操作的方式。

例如，这里我们为 x = 1, 1.1, ..., 1.8 计算 f(x) = x^2 + sqrt(x) 并计算每个点的导数（即 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]

Answer 2

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

"I'm wondering if it is possible use the autograd module (or, in general, any other autodifferentiation module) to perform this action."

答案是否定的：一旦你在space或时间中离散化你的问题，那么时间和space就变成了带有网格的离散变量-like 结构，并且不是您输入某些函数以计算 PDE 解的显式变量。

例如，如果我想计算某些流体流 u(x,t) 的速度场，我将在 space 和时间中离散化，并且我将 u[:,:] 其中indices 表示 space 和时间中的位置。

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

据我所知，您编写的张量兼容函数可能是您的最佳选择。可以看到在 PyTorch 论坛 here and here 中有人提出了类似的问题。或者你可以做类似

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

如果您不担心端点。

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

Using automatic differentiation libraries to compute partial derivatives of an arbitrary tensor

python

numpy

pde

pytorch

autograd