dot(A,B,3)的Numpy等价函数是什么？

Question

假设我有两个三维矩阵，如下所示（摘自Matlab示例http://www.mathworks.com/help/matlab/ref/dot.html）：

A = cat(3,[1 1;1 1],[2 3;4 5],[6 7;8 9])
B = cat(3,[2 2;2 2],[10 11;12 13],[14 15; 16 17])

如果我想在第三个维度上进行成对点积，我可以在matlab中这样做：

C = dot(A,B,3)

这将得到以下结果：

C =
  106   140
  178   220

有没有numpy中的等效操作，最好是矢量化选项，避免编写双重循环处理整个数组。我无法理解 np.tensordot 或 np.inner 的作用，但它们可能是选择。

- Stankalank

Answer 1

In [169]:

A = np.dstack([[[1, 1],[1 ,1]],[[2 ,3],[4, 5]],[[6, 7],[8, 9]]])
B = np.dstack([[[2, 2],[2, 2]],[[10, 11],[12, 13]],[[14, 15], [16, 17]]])
c=np.tensordot(A, B.T,1)
np.vstack([np.diag(c[:,i,i]) for i in range(A.shape[0])]).T
Out[169]:
array([[106, 140],
       [178, 220]])

但出乎意料的是，它是最慢的：

In [170]:

%%timeit
c=np.tensordot(A, B.T,1)
np.vstack([np.diag(c[:,i,i]) for i in range(A.shape[0])]).T
10000 loops, best of 3: 95.2 µs per loop
In [171]:

%timeit np.einsum('i...,i...',a,b)
100000 loops, best of 3: 6.93 µs per loop
In [172]:

%timeit inner1d(A,B)
100000 loops, best of 3: 4.51 µs per loop

- CT Zhu

---

对于两个 n x n 矩阵，一般情况是什么？ - Stankalank

出于好奇，你是如何初始化 A,B 的？ - gg349

@flebool，基本上和joshAdel的另一个答案一样，请参见编辑。 - CT Zhu

请注意，此处的 A,B 大小不正确。在 Matlab 代码中，数组的形状为 2,2,3，而您的形状为 3,2,2，需要更改程序的所有逻辑。 - gg349

@flebool，已更新。令人惊讶的是，tensordot 是最慢的。我实际上避免使用 Einstein 求和，因为我认为它会很慢。 - CT Zhu

我猜这可能是由于for循环和对diag的调用导致速度变慢。显然，我本来期望inner1d更快，但它甚至没有文档，或者至少我找不到它的文档。 - gg349

Answer 2

这里有一个解决方案：

A = dstack([[[1, 1],[1 ,1]],[[2 ,3],[4, 5]],[[6, 7],[8, 9]]])
B = dstack([[[2, 2],[2, 2]],[[10, 11],[12, 13]],[[14, 15], [16, 17]]])

C = einsum('...k,...k',A,B)

基本上，dstack 沿着第三个轴进行连接（文档），然后您可以使用 numpy 提供的强大的 Einstein 求和工具 einsum （文档）。

- gg349

---

那么假设我已经有了两个 n x n x k 的三维数组，我只需要执行 C = einsum('...k,...k',A,B) 就可以了吗？这比双重循环更有效率吗？ - Stankalank

是的，如果你习惯使用爱因斯坦求和符号，那么使用它会更高效、更易读。 - gg349

Answer 3

使用np.einsum：

In [9]: B = np.array([[[2, 2],[2, 2]],[[10, 11],[12, 13]],[[14, 15],[16, 17]]])

In [10]: A = np.array([[[1, 1],[1, 1]],[[2, 3],[4, 5]],[[6, 7],[8, 9]]]) 

In [11]: np.einsum('i...,i...',A,B)
Out[11]:
array([[106, 140],
       [178, 220]])

还有一个有趣的例子：

In [37]: from numpy.core.umath_tests import inner1d

In [38]: inner1d(A,B)
Out[38]:
array([[106, 140],
       [178, 220]])

编辑 回应 @flebool 的评论，inner1d 可以适用于形状为 (2,2,3) 和 (3,2,2) 的数组：

In [41]: A = dstack([[[1, 1],[1 ,1]],[[2 ,3],[4, 5]],[[6, 7],[8, 9]]])

In [42]: B = dstack([[[2, 2],[2, 2]],[[10, 11],[12, 13]],[[14, 15], [16, 17]]])

In [43]: inner1d(A,B)
Out[43]:
array([[106, 140],
       [178, 220]])

- JoshAdel

---

请注意，此处的 A,B 大小不正确。在 Matlab 代码中，数组的形状为 2,2,3，而您的形状为 3,2,2，需要更改程序的所有逻辑。 - gg349

这似乎是最干净的解决方案。与einsum相比，它是否有任何性能缺陷？ - Stankalank

@Stankalank，针对您在原帖中的测试数组，einsum=5.33微秒每次循环，inner1d=4.3微秒每次循环。对于一个(2,2,10000)的数组，einsum=42.6微秒每次循环，inner1d=41.6微秒每次循环，但是计时会因硬件和numpy版本/编译设置而有所不同。 - JoshAdel

inner1d 的文档在哪里？它的 __doc__ 内容相当简略。 - hpaulj

@hpaulj 我记不得我第一次是在哪里发现这种方法的，但你可以从导入中看出它并不是一个显而易见的方法。你可以查看源代码（C 代码）：https://github.com/numpy/numpy/blob/master/numpy/core/src/umath/umath_tests.c.src - JoshAdel