Question

我试图弄清楚如何在10x10x10 3D网格上广播系统点产品操作。我制定了以下数组：

A.shape=(5,5,10,10,10)
b.shape=(5,10,10,10)

我想获得如下的数组

c.shape=(5,10,10,10)

到目前为止，我通过以下代码获得了

c=np.sum(A*b,axis=1)

但是，我相信我应该能够使用np.dot或np.tensordot获得相同的结果。我努力了，但我无法获得相同的结果。这样做有助于我了解np.tensordot的工作原理，因为我还需要在我的工作中使用np.tensorsolve。

Answer 1

我们需要保持最后四个轴对齐并且除了第二个轴（轴= 1）之外还要在输出中保持这些轴，这将是总和减少的。对于这种情况，除了np.matmul之外，np.einsum是可行的方法。使用np.einsum，它会更容易/更直观 -

c = np.einsum('ijklm,jklm->iklm',A,b)

Answer 2

如果＆＃39;点＆＃39;尺寸在最后，matmul将起作用

比较3种方法：

In [252]: A=np.arange(5*5*10*10*10).reshape((5,5,10,10,10))
In [253]: B=np.arange(5*10*10*10).reshape((5,10,10,10))

In [254]: C=np.sum(A*B, axis=1)
In [255]: D=np.einsum('ijklm,jklm->iklm',A,B)
In [256]: E = (A.transpose(2,3,4,0,1)@B.transpose(1,2,3,0)[...,None])[...,0].transpose(3,0,1,2)

所有转置都会使数组进入(....,5,5)和(...,5,1)。

In [257]: np.allclose(C,D)
Out[257]: True
In [258]: np.allclose(C,E)
Out[258]: True

In [259]: timeit C=np.sum(A*B, axis=1)
124 µs ± 4 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
In [260]: timeit D=np.einsum('ijklm,jklm->iklm',A,B)
66 µs ± 18.7 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
In [261]: timeit E = (A.transpose(2,3,4,0,1)@B.transpose(1,2,3,0)[...,None])[...
     ...: ,0].transpose(3,0,1,2)
68.6 µs ± 973 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

tensordot重新整形并转置数组，以便它可以执行简单的2d dot - 然后转换回来。最后3个维度实际上是单个1000维度。

In [262]: np.tensordot(A,B,(1,0)).shape
Out[262]: (5, 10, 10, 10, 10, 10, 10)
In [263]: timeit np.tensordot(A,B,(1,0)).shape
74 ms ± 70.1 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

结果更大 - 一种非求和维度的外部产品。结果在那里，埋藏为多维对角线。

tensordot实际上是：

In [264]: X=(B.reshape(5,1000).T).dot(A.reshape(5,5,1000))
In [265]: X.shape
Out[265]: (1000, 5, 1000)
In [266]: timeit X=(B.reshape(5,1000).T).dot(A.reshape(5,5,1000))
78.9 ms ± 82.6 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

使用numpy

2 个答案: