Python中高维度的Numpy矩阵乘法

Question

Python中高维度的Numpy矩阵乘法

pythonnumpymatrixmachine-learninglinear-algebra

3

我正在尝试寻找一个在numpy中能够加速以下计算的矩阵运算。

我有两个3D矩阵A和B。第一维表示样例，在两个矩阵中都有n_examples个例子。我想要实现的是对A和B中的每个示例进行点积并求和的结果：

import numpy as np

n_examples = 10
A = np.random.randn(n_examples, 20,30)
B = np.random.randn(n_examples, 30,5)
sum = np.zeros([20,5])
for i in range(len(A)):
  sum += np.dot(A[i],B[i])

- aha

2个回答

2

哈，只需要一行代码就能完成：np.einsum('nmk,nkj->mj',A,B)。

参见爱因斯坦求和约定：http://docs.scipy.org/doc/numpy/reference/generated/numpy.einsum.html

虽然问题不同，但思路基本相同。在我们刚刚讨论的这个话题中，可以看到讨论和替代方法：numpy multiply matrices preserve third axis

不要把变量命名为sum，否则会覆盖内置的sum函数。

正如@Jaime所指出的，对于这些大小的维度，循环实际上更快。事实上，基于map和sum的解决方案，虽然更简单，但速度更慢：

In [19]:

%%timeit
SUM = np.zeros([20,5])
for i in range(len(A)):
  SUM += np.dot(A[i],B[i])
10000 loops, best of 3: 115 µs per loop
In [20]:

%timeit np.array(map(np.dot, A,B)).sum(0)
1000 loops, best of 3: 445 µs per loop
In [21]:

%timeit np.einsum('nmk,nkj->mj',A,B)
1000 loops, best of 3: 259 µs per loop

当维度变得更加庞大时，情况会有所不同：

n_examples = 1000
A = np.random.randn(n_examples, 20,1000)
B = np.random.randn(n_examples, 1000,5)

并且：

In [46]:

%%timeit
SUM = np.zeros([20,5])
for i in range(len(A)):
  SUM += np.dot(A[i],B[i])
1 loops, best of 3: 191 ms per loop
In [47]:

%timeit np.array(map(np.dot, A,B)).sum(0)
1 loops, best of 3: 164 ms per loop
In [48]:

%timeit np.einsum('nmk,nkj->mj',A,B)
1 loops, best of 3: 451 ms per loop

- CT Zhu

1

它比OP的代码慢50％，对于非常大的输入则更糟糕。 - Jaime

一种针对更大维度略微更快的方法，但并没有令人印象深刻的改进。einsum 变得更慢了。现在必须睡觉了，希望能从西海岸得到解决方案。 :P - CT Zhu

它比tensordot慢吗？ - MartianMartian

你是否知道其他编程语言的库是否提供了更优化的解决方案？ - MartianMartian

我建议你使用timeit测试你的代码，只要确保你所做的计算确实比其他方法慢。请参考相关讨论：https://dev59.com/FmMl5IYBdhLWcg3we3Do. - CT Zhu

此外，numpy 版本是用 C 编写的 https://github.com/numpy/numpy/blob/master/numpy/core/src/multiarray/einsum.c.src，您也可以向作者寻求建议。 - CT Zhu

网页内容由stack overflow 提供, 点击上面的

可以查看英文原文，
原文链接

- Saullo G. P. Castro · Accepted Answer

这是np.tensordot()的一个典型应用：

sum = np.tensordot(A, B, [[0,2],[0,1]])

时间

使用以下代码：

import numpy as np

n_examples = 100
A = np.random.randn(n_examples, 20,30)
B = np.random.randn(n_examples, 30,5)

def sol1():
    sum = np.zeros([20,5])
    for i in range(len(A)):
      sum += np.dot(A[i],B[i])
    return sum

def sol2():
    return np.array(map(np.dot, A,B)).sum(0)

def sol3():
    return np.einsum('nmk,nkj->mj',A,B)

def sol4():
    return np.tensordot(A, B, [[2,0],[1,0]])

def sol5():
    return np.tensordot(A, B, [[0,2],[0,1]])

结果：

timeit sol1()
1000 loops, best of 3: 1.46 ms per loop

timeit sol2()
100 loops, best of 3: 4.22 ms per loop

timeit sol3()
1000 loops, best of 3: 1.87 ms per loop

timeit sol4()
10000 loops, best of 3: 205 µs per loop

timeit sol5()
10000 loops, best of 3: 172 µs per loop

在我的电脑上，tensordot() 是最快的解决方案，改变轴被评估的顺序既不会改变结果也不会影响性能。