如何在 autograd 中逐个元素分配方程式

Question

我正在尝试为非线性 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 中分配元素。这是真的吗。有什么解决方法吗？或者可以推荐另一个图书馆？

谢谢！

Answer 1

确实，数组索引赋值在 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 and 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.]

"Or maybe another library to suggest?"

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

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

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

如何在 autograd 中逐个元素分配方程式

How to assign equations element by element in autograd

python

automatic-differentiation

autograd