Как поэлементно назначать уравнения в автограде

Я пытаюсь реализовать решатель на основе автограда для нелинейного УЧП. Как и в случае с большинством PDE, мне нужно иметь возможность работать с отдельными элементами моего входного вектора, но, по-видимому, это нарушает автоград. Я создал этот простой пример, чтобы показать проблему, с которой я столкнулся:

Работает следующий код:

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)

Я читал в нескольких местах, что вы не можете назначать элементы в автограде. Действительно ли это так. Есть ли обходной путь? Или может другую библиотеку предложить?

Благодарю вас!


python autograd automatic-differentiation
person Rafael March    schedule 10.12.2019    source источник

Ответы (1)


arrow_upward
0
arrow_downward

Действительно, назначение индексации массива невозможно в autograd. Люди написали решатели PDE в autograd (см. 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.]

"Или, может быть, другую библиотеку посоветуете?"

Взгляните на дельфин-прилегающий. Если вы можете написать свой решатель PDE в FEniCS, это может быть полезно.

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

Обратите внимание, что наивное обратное распространение через средство решения нелинейных уравнений может быть неэффективным способом вычисления производных скалярной функции потерь, см. https://cs.stanford.edu/~ambrad/adjoint_tutorial.pdf за довольно основательное руководство по сопряженному методу, в том числе для нелинейных задач.

person Nick McGreivy    schedule 24.01.2020