高維中的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()是最快的解決方案，並且更改軸的評估順序並沒有改變結果和性能。

Answer 2

哈，它可以在一行中完成： np.einsum('nmk,nkj->mj',A,B) 。

見愛因斯坦總結： http ： //docs.scipy.org/doc/numpy/reference/generated/numpy.einsum.html

不一樣的問題但是這個想法是完全相同的，請參閱我們剛剛討論過的這個主題中的討論和替代方法： numpy multiply矩陣保留第三軸

不要命名變量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