介绍和解决方案代码
np.einsum
非常强大,但在极少数情况下,如果你可以将矩阵乘法
引入计算中,仍然可以超越它。经过几次尝试,似乎可以使用np.dot进行矩阵乘法
来提高性能,而不是使用np.einsum('ia,aj,ka->ijk', A, B, C)
。
基本思想是将“所有einsum”操作分解成以下np.einsum
和np.dot
的组合:
- 对于
A:[i,a]
和B:[a,j]
的求和使用np.einsum
得到一个3D数组:[i,j,a]
。
- 然后将此3D数组重塑为
2D数组:[i*j,a]
,并且第三个数组C[k,a]
被转置为[a,k]
,以便在这两个数组之间执行矩阵乘法
,从而得到[i*j,k]
作为矩阵乘积,因为我们在那里失去了索引[a]
。
- 将乘积重塑为一个
3D数组:[i,j,k]
,作为最终输出。
以下是迄今为止讨论的第一个版本的实现 -
import numpy as np
def tensor_prod_v1(A,B,C):
m,d = A.shape
n = B.shape[1]
p = C.shape[0]
AB = np.einsum('ia,aj->ija', A, B)
return np.dot(AB.reshape(m*n,d),C.T).reshape(m,n,p)
由于我们要对三个输入数组的a-th
索引进行求和,因此可以有三种不同的方法沿着a-th
索引进行求和。之前列出的代码是针对(A,B)
的。因此,我们还可以有(A,C)
和(B,C)
,这给我们带来了另外两种变化,如下所示:
def tensor_prod_v2(A,B,C):
m,d = A.shape
n = B.shape[1]
p = C.shape[0]
AC = np.einsum('ia,ja->ija', A, C)
return np.dot(AC.reshape(m*p,d),B).reshape(m,p,n).transpose(0,2,1)
def tensor_prod_v3(A,B,C):
m,d = A.shape
n = B.shape[1]
p = C.shape[0]
BC = np.einsum('ai,ja->aij', B, C)
return np.dot(A,BC.reshape(d,n*p)).reshape(m,n,p)
根据输入数组的形状,不同的方法会相对于彼此产生不同的加速效果,但我们希望所有方法都比all-einsum
方法更好。性能数据列在下一节中。
运行时测试
这可能是最重要的部分,因为我们试图查看所提出的三种方法与问题中最初提出的all-einsum
方法相比的加速数值。
数据集#1(等形状的数组):
In [494]: L1 = 200
...: L2 = 200
...: L3 = 200
...: al = 200
...:
...: A = np.random.rand(L1,al)
...: B = np.random.rand(al,L2)
...: C = np.random.rand(L3,al)
...:
In [495]: %timeit tensor_prod_v1(A,B,C)
...: %timeit tensor_prod_v2(A,B,C)
...: %timeit tensor_prod_v3(A,B,C)
...: %timeit np.einsum('ia,aj,ka->ijk', A, B, C)
...:
1 loops, best of 3: 470 ms per loop
1 loops, best of 3: 391 ms per loop
1 loops, best of 3: 446 ms per loop
1 loops, best of 3: 3.59 s per loop
数据集 #2(更大的A):
In [497]: L1 = 1000
...: L2 = 100
...: L3 = 100
...: al = 100
...:
...: A = np.random.rand(L1,al)
...: B = np.random.rand(al,L2)
...: C = np.random.rand(L3,al)
...:
In [498]: %timeit tensor_prod_v1(A,B,C)
...: %timeit tensor_prod_v2(A,B,C)
...: %timeit tensor_prod_v3(A,B,C)
...: %timeit np.einsum('ia,aj,ka->ijk', A, B, C)
...:
1 loops, best of 3: 442 ms per loop
1 loops, best of 3: 355 ms per loop
1 loops, best of 3: 303 ms per loop
1 loops, best of 3: 2.42 s per loop
数据集 #3 (更大的 B) :
In [500]: L1 = 100
...: L2 = 1000
...: L3 = 100
...: al = 100
...:
...: A = np.random.rand(L1,al)
...: B = np.random.rand(al,L2)
...: C = np.random.rand(L3,al)
...:
In [501]: %timeit tensor_prod_v1(A,B,C)
...: %timeit tensor_prod_v2(A,B,C)
...: %timeit tensor_prod_v3(A,B,C)
...: %timeit np.einsum('ia,aj,ka->ijk', A, B, C)
...:
1 loops, best of 3: 474 ms per loop
1 loops, best of 3: 247 ms per loop
1 loops, best of 3: 439 ms per loop
1 loops, best of 3: 2.26 s per loop
数据集#4(更大的C):
In [503]: L1 = 100
...: L2 = 100
...: L3 = 1000
...: al = 100
...:
...: A = np.random.rand(L1,al)
...: B = np.random.rand(al,L2)
...: C = np.random.rand(L3,al)
In [504]: %timeit tensor_prod_v1(A,B,C)
...: %timeit tensor_prod_v2(A,B,C)
...: %timeit tensor_prod_v3(A,B,C)
...: %timeit np.einsum('ia,aj,ka->ijk', A, B, C)
...:
1 loops, best of 3: 250 ms per loop
1 loops, best of 3: 358 ms per loop
1 loops, best of 3: 362 ms per loop
1 loops, best of 3: 2.46 s per loop
数据集 #5(更长的a维度长度):
In [506]: L1 = 100
...: L2 = 100
...: L3 = 100
...: al = 1000
...:
...: A = np.random.rand(L1,al)
...: B = np.random.rand(al,L2)
...: C = np.random.rand(L3,al)
...:
In [507]: %timeit tensor_prod_v1(A,B,C)
...: %timeit tensor_prod_v2(A,B,C)
...: %timeit tensor_prod_v3(A,B,C)
...: %timeit np.einsum('ia,aj,ka->ijk', A, B, C)
...:
1 loops, best of 3: 373 ms per loop
1 loops, best of 3: 269 ms per loop
1 loops, best of 3: 299 ms per loop
1 loops, best of 3: 2.38 s per loop
结论:我们发现,相比于问题中列出的 all-einsum
方法,采用提出的方法的变体可以实现 8x-10x
的加速。
bsxfun
和一个矩阵乘法来实现这个功能,我觉得这两个操作都很快速。这样对你来说可以吗? - Luis Mendo