高维Python Numpy矩阵乘法

AHa 发表于 Dev

啊哈

我正在尝试寻找一种矩阵运算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])

索洛GP卡斯特罗

这是典型的应用程序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()是最快的解决方案，改变轴的评估顺序不会改变结果，也不会改变性能。

本文收集自互联网，转载请注明来源。

如有侵权，请联系[email protected] 删除。