Python 用numpy乘高阶矩阵

我制造了这个玩具问题，它反映了我更大的问题：

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应该是最好的方法，因为这个操作很重，而且我的矩阵很大

我看到了，但是形状是不同的，我不能理解我应该使用哪个轴来解决这个问题。然而，我相信这应该有答案。有什么建议吗

---

编辑：我接受了，但也请阅读我自己的答案，它可能会帮助其他人…

您可以使用

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

第一轴中的条目与

dist

第一轴中的条目相乘（下标

'i'

）。每个数组的第二个轴同上（下标

'j'

）。我们不返回3D数组，而是告诉einsum通过从输出下标中省略它来沿轴

'j'

求和值，从而返回2D数组

---

np.tensordot

更难应用于这个问题。它会自动对轴的乘积求和。然而，我们想要两套产品，但只求其中一套的和

编写

np.tensordot（ind，dist，axes=[1，1]）

（如您链接的答案中所示）会为您计算正确的值，但会返回一个具有形状

（3，4，3）

的3D数组。如果您能够负担更大阵列的内存成本，则可以使用：

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

这会给出正确的结果，但是因为

tensordot

首先创建了一个比必需的数组大得多的数组，

einsum

似乎是一个更好的选择。

接下来，我想确定哪种方法最快，所以我使用了

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=（3000200400）

，翻转顺序以确保其不起作用，并设置了另一个重复次数较多的测试：

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

小得多，所以我选择了丑陋的版本，在这个例子中也是最快的。

有趣的是，

einsum

并不一定比在循环中使用

dot

快！我认为，如果第一个维度很大，那么

for

循环将

dot

向下拖动很多，而

einsum

可能是更快的选择。我尝试了

I，j，k=（300000,20,40）

，得到

dot

为1.11秒，而

einsum

为273毫秒。正如你的回答所示，在你正在使用的形状上测试这些方法绝对值得，这样你就知道什么是最快的了。@ajcr我完全同意！我也会添加这些结果（编辑：我接受了@ajcr的答案，但也请阅读我自己的答案。）这将对其他人有所帮助。。