在高维Python Numpy矩阵乘法

Question

我想在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]) 

Answer 1

这是一个典型应用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()是最快的解决方案，改变为使评估轴并不会改变结果n性能。

Answer 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