将高阶矩阵与 numpy 相乘

Posted 2023-03-12

【中文标题】将高阶矩阵与 numpy 相乘

【英文标题】：Multiply high order matrices with numpy

【发布时间】：2015-08-29 06:30:18

【问题描述】：

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

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

我看到了this post，但形状不同，我不明白应该使用哪个axes 来解决这个问题。但是，我确信tensordot 应该有答案。有什么建议吗？

编辑：我接受了@ajcr's answer，但也请阅读我自己的答案，它可能对其他人有所帮助...

【问题讨论】：

太棒了（编辑：我接受了@ajcr 的回答，但也请阅读我自己的回答。）它将对其他人有所帮助..

【参考方案1】：

在@ajcr's great answer之后，我想确定哪种方法最快，所以我使用了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 是错误的做法（更不用说在更大的例子中memory error，正如@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

有关这些结果的解释，请参见 cmets。

道德是，在为特定问题寻找最快的解决方案时，尝试生成在类型和形状方面尽可能与原始数据相似的数据。在我的情况下，i 比 j,k 小得多，所以我选择了丑陋的版本，这也是这种情况下最快的版本。

【讨论】：

有趣的是einsum 并不一定比在循环中使用dot 更快！我认为如果第一个维度很大，for 循环会将dot 拖下很多，einsum 可能是更快的选择。我尝试使用 i,j,k = (300000,20,40) 并得到 dot 为 1.11 秒，而 einsum 为 273 毫秒。正如您的回答所示，绝对值得在您正在使用的形状上测试这些方法，这样您就知道什么是最快的。

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

【参考方案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 的第一个轴中的条目与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 似乎是一个更好的选择。