使用numpy计算高阶矩阵的乘积

Question

我创建了这个玩具问题，以反映我更大的问题：

import numpy as np

ind = np.ones((3,2,4)) # shape=(3L, 2L, 4L)
dist = np.array([[0.1,0.3],[1,2],[0,1]]) # shape=(3L, 2L)

ans = np.array([np.dot(dist[i],ind[i]) for i in xrange(dist.shape[0])]) # shape=(3L, 4L)
print ans

""" prints:
   [[ 0.4  0.4  0.4  0.4]
    [ 3.   3.   3.   3. ]
    [ 1.   1.   1.   1. ]]
"""

为了尽可能快地完成它，使用Numpy的函数来计算ans应该是最好的方法，因为这个操作很重要，而且我的矩阵很大。

我看到这篇文章，但形状不同，我无法理解应该为这个问题使用哪些axes。然而，我确定tensordot应该有答案。任何建议？

编辑：我接受了@ajcr's answer，但请阅读我的答案，它可能会帮助其他人……

- omerbp

---

太好了，这将对其他人有所帮助。 - Imran Ali Khan

Answer 1

在参考@ajcr的优秀回答后，我想确定哪种方法最快，所以我使用了timeit：

import timeit

setup_code = """
import numpy as np
i,j,k = (300,200,400)
ind = np.ones((i,j,k)) #shape=(3L, 2L, 4L)
dist = np.random.rand(i,j) #shape=(3L, 2L)
"""

basic ="np.array([np.dot(dist[l],ind[l]) for l in xrange(dist.shape[0])])"
einsum = "np.einsum('ijk,ij->ik', ind, dist)"
tensor= "np.tensordot(ind, dist, axes=[1, 1])[0].T"

print "tensor - total time:", min(timeit.repeat(stmt=tensor,setup=setup_code,number=10,repeat=3))
print "basic - total time:", min(timeit.repeat(stmt=basic,setup=setup_code,number=10,repeat=3))
print "einsum - total time:", min(timeit.repeat(stmt=einsum,setup=setup_code,number=10,repeat=3))

意外的结果是：

tensor - total time: 6.59519493952
basic - total time: 0.159871203461
einsum - total time: 0.263569731028

很明显，使用tensordot是错误的方法（正如@ajcr所述，在更大的示例中还会出现“内存错误”）。由于这个示例很小，我将矩阵的大小改为 i，j，k =（3000,200,400），反转顺序只是为了确保它没有影响，并设置了另一个测试，重复次数更高：

print "einsum - total time:", min(timeit.repeat(stmt=einsum,setup=setup_code,number=50,repeat=3))
print "basic - total time:", min(timeit.repeat(stmt=basic,setup=setup_code,number=50,repeat=3))

结果与第一次运行一致：

einsum - total time: 13.3184077671
basic - total time: 8.44810031351

然而，测试另一种尺寸增长的方法 - i,j,k = (30000,20,40) - 得到了以下结果：

einsum - total time: 0.325594117768
basic - total time: 0.926416766397

请查看注释以了解这些结果的解释。

结论是，在寻找特定问题的最快解决方案时，请尽量生成与原始数据类型和形状相似的数据。在我的情况下，i比j,k小得多，因此我选择了丑陋的版本，这也是本例中最快的版本。

- omerbp

---

有趣的是，使用循环中的dot并不一定比einsum快！如果第一个维度很大，那么for循环会拖慢dot的速度，这时候einsum可能会成为更快的选择。我尝试使用i,j,k = (300000,20,40)，得到使用dot需要1.11秒，而使用einsum只需要273毫秒。就像你的回答所展示的那样，在处理自己工作中使用的形状时测试这些方法肯定是值得的，这样你就能知道哪个方法最快了。 - Alex Riley

@ajcr 我完全同意！我也会添加这些结果。 - AvidLearner

Answer 2

您可以使用np.einsum来执行该操作，因为它允许对要乘的轴和要相加的轴进行非常精细的控制：

>>> np.einsum('ijk,ij->ik', ind, dist)
array([[ 0.4,  0.4,  0.4,  0.4],
       [ 3. ,  3. ,  3. ,  3. ],
       [ 1. ,  1. ,  1. ,  1. ]])

该函数将ind数组第一轴(下标为'i')的每个元素与 dist数组第一轴(下标为'i')的对应元素进行相乘。同样，将这两个数组的第二轴(下标为'j')的每个元素相乘。通过在输出下标中省略轴'j'，einsum返回一个沿着轴'j'求和的2D数组而不是3D数组。

---

np.tensordot在解决这个问题时会更加困难。它会自动对轴进行求和。然而，在此我们希望得到两组乘积结果，但只对其中一组求和。

使用np.tensordot(ind, dist, axes=[1, 1])可以计算出正确的值，但它返回一个形状为(3,4,3)的3D数组。如果你可以承担更大的内存成本，你可以使用：

np.tensordot(ind, dist, axes=[1, 1])[0].T

这样可以得到正确的结果，但由于tensordot事先创建了一个比必要大得多的数组，因此einsum似乎是更好的选择。

- Alex Riley