`numpy.einsum`中的`out`参数无法按预期工作

时间:2017-11-26 01:56:52

标签: numpy python-3.5

我有两个代码。第一个是:

A = np.arange(3*4*3).reshape(3, 4, 3)
P = np.arange(1, 4)
A[:, 1:, :] = np.einsum('j, ijk->ijk', P, A[:, 1:, :])

,结果A为:

array([[[  0,   1,   2],
        [  6,   8,  10],
        [ 18,  21,  24],
        [ 36,  40,  44]],

       [[ 12,  13,  14],
        [ 30,  32,  34],
        [ 54,  57,  60],
        [ 84,  88,  92]],

       [[ 24,  25,  26],
        [ 54,  56,  58],
        [ 90,  93,  96],
        [132, 136, 140]]])

第二个是:

A = np.arange(3*4*3).reshape(3, 4, 3)
P = np.arange(1, 4)
np.einsum('j, ijk->ijk', P, A[:, 1:, :], out=A[:,1:,:])

,结果A为:

array([[[ 0,  1,  2],
        [ 0,  0,  0],
        [ 0,  0,  0],
        [ 0,  0,  0]],

       [[12, 13, 14],
        [ 0,  0,  0],
        [ 0,  0,  0],
        [ 0,  0,  0]],

       [[24, 25, 26],
        [ 0,  0,  0],
        [ 0,  0,  0],
        [ 0,  0,  0]]])

所以结果是不同的。在这里,我想使用out来节省内存。这是numpy.einsum中的错误吗?或者我错过了什么?

顺便说一句,我的numpy版本是1.13.3。

2 个答案:

答案 0 :(得分:4)

我以前没有使用过这个新的out参数,但过去曾与einsum合作,并且大致了解它是如何工作的(或者至少习惯了)。

我认为它在迭代开始之前将out数组初始化为零。这将考虑A[:,1:,:]块中的所有0。如果我初始分开out数组,则插入所需的值

In [471]: B = np.ones((3,4,3),int)
In [472]: np.einsum('j, ijk->ijk', P, A[:, 1:, :], out=B[:,1:,:])
Out[472]: 
array([[[  3,   4,   5],
        [ 12,  14,  16],
        [ 27,  30,  33]],

       [[ 15,  16,  17],
        [ 36,  38,  40],
        [ 63,  66,  69]],

       [[ 27,  28,  29],
        [ 60,  62,  64],
        [ 99, 102, 105]]])
In [473]: B
Out[473]: 
array([[[  1,   1,   1],
        [  3,   4,   5],
        [ 12,  14,  16],
        [ 27,  30,  33]],

       [[  1,   1,   1],
        [ 15,  16,  17],
        [ 36,  38,  40],
        [ 63,  66,  69]],

       [[  1,   1,   1],
        [ 27,  28,  29],
        [ 60,  62,  64],
        [ 99, 102, 105]]])

einsum的Python部分并没有告诉我什么,除了它如何决定将out数组传递给c部分,(作为其中一个列表tmp_operands):

c_einsum(einsum_str,* tmp_operands,** einsum_kwargs)

我知道它设置c-api等效于np.nditer,使用str来定义轴和迭代。

它迭代迭代教程中的这一部分:

https://docs.scipy.org/doc/numpy-1.13.0/reference/arrays.nditer.html#reduction-iteration

特别注意it.reset()步骤。这会在迭代之前将out缓冲区设置为0。然后迭代输入数组和输出数组的元素,将计算值写入输出元素。由于它正在进行一系列产品(例如out[:] += ...),因此必须从一个干净的平板开始。

我猜测实际发生了什么,但我认为它应该将输出缓冲区清零以开始。如果该数组与其中一个输入相同,那么最终会搞乱计算。

所以我认为这种方法不会起作用并且可以节省你的记忆力。它需要一个干净的缓冲区来累积结果。一旦完成它,或者你,可以将值写回A。但考虑到dot类似产品的性质,您不能使用相同的数组进行输入和输出。

In [476]: A[:,1:,:] = np.einsum('j, ijk->ijk', P, A[:, 1:, :])
In [477]: A
Out[477]: 
array([[[  0,   1,   2],
        [  3,   4,   5],
        [ 12,  14,  16],
        [ 27,  30,  33]],
        ....)

答案 1 :(得分:3)

einsum的C源代码中,there is a section将采用out指定的数组并进行零设置。

但是在Python source code中,有一些执行路径在调用tensordot的参数之前调用c_einsum函数。

这意味着在任何子数组被设置之前,可能会预先计算某些操作(因此在某些收缩过程中修改数组Atensordot einsum的C代码中的零设置器为零。

另一种说法是:在进行下一次收缩操作的每次传递中,NumPy有很多选择。要直接使用tensordot而不进入C级einsum代码呢?或者准备参数并传递给C级(这将涉及用全零覆盖输出数组的一些子视图)?或者重新订购操作并重复检查?

根据它为这些优化选择的顺序,您最终可能会得到意外的全零子阵列。

最好的办法是不要试图变得聪明并使用相同的数组作为输出。你说这是因为你想节省内存。是的,在某些特殊情况下,einsum操作可能就地可行。但它目前没有检测到是否是这种情况并试图避免置零。

在很多情况下,在整个操作过程中过度写入其中一个输入数组会导致很多问题,就像尝试附加到直接循环的列表一样,等等。