0

我有一个相当重要的函数,我想区分它,autograd但我还不够一个 numpy 向导来弄清楚我们如何在没有数组分配的情况下做到这一点。

我也很抱歉,为了能够独立运行,我不得不让这个例子难以置信地做作和毫无意义。我正在使用的实际代码是针对非线性有限元的,并且正在尝试为复杂的非线性系统计算雅可比。

import autograd.numpy as anp
from autograd import jacobian

def alpha(x):
    return anp.exp(-(x - 10) ** 2) / (x + 1)


def f(x):
    # Matrix getting constructed
    k = anp.zeros((x.shape[0], x.shape[0]))

    # loop over some random 3 dimensional vectors
    for element in anp.random.randint(0, x.shape[0], (x.shape[0], 3)):

        # select 3 values from x
        x_ijk = anp.array([[x[i] for i in element]])

        norm = anp.linalg.norm(
            x_ijk @ anp.vstack((element, element)).transpose()
        )

        # make some matrix from the element
        m = element.reshape(3, 1) @ element.reshape(1, 3)

        # alpha is an arbitrary differentiable function R -> R
        alpha_value = alpha(norm)

        # combine m matricies into k scaling by alpha_value
        n = m.shape[0]
        for i in range(n):
            for j in range(n):
                k[element[i], element[j]] += m[i, j] * alpha_value

    return k @ x


print(jacobian(f)(anp.random.rand(10)))
# And course we get an error
# k[element[i], element[j]] += m[i, j] * alpha_value
# ValueError: setting an array element with a sequence.

我不太明白这条消息,因为没有发生类型错误。我认为它必须来自分配。

写完上面的内容后,我做了一个简单的切换,PyTorch代码运行得很好。但我仍然更喜欢使用autograd

#pytorch version
import torch
from torch.autograd.gradcheck import zero_gradients


def alpha(x):
    return torch.exp(x)


def f(x):
    # Matrix getting constructed
    k = torch.zeros((x.shape[0], x.shape[0]))

    # loop over some random 3 dimensional vectors
    for element in torch.randint(0, x.shape[0], (x.shape[0], 3)):

        # select 3 values from x
        x_ijk = torch.tensor([[1. if n == e else 0 for n in range(len(x))] for e in element]) @ x
        norm = torch.norm(
            x_ijk @ torch.stack((torch.tanh(element.float() + 4), element.float() - 4)).t()
        )

        m = torch.rand(3, 3)

        # alpha is an arbitrary differentiable function R -> R
        alpha_value = alpha(norm)

        n = m.shape[0]
        for i in range(n):
            for j in range(n):
                k[element[i], element[j]] += m[i, j] * alpha_value

    print(k)
    return k @ x


x = torch.rand(4, requires_grad=True)
print(x, '\n')
y = f(x)
print(y, '\n')
grads = []
for val in y:
    val.backward(retain_graph=True)
    grads.append(x.grad.clone())
    zero_gradients(x)
if __name__ == '__main__':
    print(torch.stack(grads))
4

1 回答 1

0

在 Autograd 和 JAX 中,不允许执行数组索引分配。有关这方面的部分解释,请参阅JAX 陷阱。

PyTorch 允许此功能。如果你想在 autograd 中运行你的代码,你必须找到一种方法来删除违规行k[element[i], element[j]] += m[i, j] * alpha_value。如果您可以在 JAX 中运行代码(其语法与 autograd 基本相同,但功能更多),那么看起来jax.ops可能有助于执行这种索引分配。

于 2020-01-24T08:31:12.557 回答