0

我正在尝试为非线性 PDE 实现基于 autograd 的求解器。与大多数 PDE 一样，我需要能够在输入向量的各个条目中进行操作，但显然这会破坏 autograd。我创建了这个简单的示例来展示我面临的问题：

以下代码有效：

def my_equation(x):
    eq = x
    return eq

x = np.random.randn(2,)
jac = autograd.jacobian(my_equation)
jacval = jac(x)
print(jacval)

以下代码不起作用：

def my_equation(x):
    eq = x
    # This breaks the code, although is a 
    # trivial line
    eq[1] = x[1]
    return eq

x = np.random.randn(2,)
jac = autograd.jacobian(my_equation)
jacval = jac(x)
print(jacval)

我在几个地方读过你不能在 autograd 中分配元素。这真的是真的吗。有什么解决方法吗？或者也许另一个图书馆建议？

谢谢！

4

1 回答 1

0

事实上，数组索引分配在 autograd 中是不可能的。人们已经在 autograd 中编写了 PDE 求解器（请参阅https://github.com/HIPS/autograd/tree/master/examples/fluidsim），所以也许有一种方法可以在留在 autograd 的同时解决您的问题。

JAX 似乎为 jax.ops 包提供了一种解决方法（请参阅https://jax.readthedocs.io/en/latest/jax.ops.html和https://github.com/google/jax#current-gotchas） .

在 PyTorch 中似乎可以进行数组索引。这表明 PyTorch 将是您的出路。以下代码有效。

import torch

def f(x):
    eq = 2*x 
    eq[0] = x[0] 
    return eq

x = torch.rand(4, requires_grad=True) 
y = f(x)
z = torch.sum(y)
z.backward()
print(x.grad) # prints [1., 2., 2., 2.]

“或者也许可以推荐另一个图书馆？”

看看 dolfin-adjoint。如果您可以在 FEniCS 中编写 PDE 求解器，那么它可能会有所帮助。

http://www.dolfin-adjoint.org/en/latest/

请注意，通过非线性方程求解器的简单反向传播可能是计算标量损失函数导数的低效方法，请参阅https://cs.stanford.edu/~ambrad/adjoint_tutorial.pdf以获得相当可靠的教程伴随方法，包括非线性问题。

于 2020-01-24T07:10:03.840 回答