4

我试图用来numpy.einsum简化我的代码中的循环。

目前，我的代码如下所示：

k = 100 
m = 50
n = 10
A = np.arange(k*m*n).reshape(k, m, n)
B = np.arange(m*m).reshape(m, m)

T = np.zeros((n, n)) 
for ind in xrange(k):
    T += np.dot(A[ind,:,:].T, np.dot(B, A[ind,:,:]))

我正在尝试numpy.einsum用作此循环的替代方案：

Tp = np.einsum('nij,njk->ik', np.einsum('nij,jk->nik', A.transpose(0,2,1), B), A)
print np.allclose(T, Tp)

是否可以使用一个numpy.einsum而不是两个？

4

2 回答 2

4

在我的电脑上，你的时间是：

np.einsum('nij,njk->ik', np.einsum('nij,jk->nik', A.transpose(0,2,1), B), A)
# 100 loops, best of 3: 4.55 ms per loop

您可以通过以下方式实现：

T2 = np.einsum('nij, il, kln ->jk', A, B, A.T)
#  10 loops, best of 3: 51.9 ms per loop

或使用双np.tensordot()：

T3 = np.tensordot(A, np.tensordot(A, B, axes=(1, 1)), axes=((0, 1), (0, 2)))
# 100 loops, best of 3: 2.73 ms per loop

我的结论是，通过两步执行此操作，您将获得更好的性能。这可能是由于一次执行操作时发生的较大步幅，可能导致更多的缓存丢失。

于 2015-10-28T20:16:25.307 回答
3

您在另一个答案中得到的真正聪明的方法。我想将numpy.dot基于方法添加到也使用一些reshaping. 这是这样做的一种方法 -

k,m,n = A.shape
((A.transpose(2,0,1).reshape(-1,m).dot(B.T)).reshape(n,-1)).dot(A.reshape(-1,n)).T

运行时测试 -

本节比较了本文前面列出的所有方法other answer和numpy.dot基于本文前面列出的方法。

In [130]: k = 100 
     ...: m = 50
     ...: n = 10
     ...: A = np.arange(k*m*n).reshape(k, m, n)
     ...: B = np.arange(m*m).reshape(m, m)
     ...: 

In [131]: %timeit np.einsum('nij,njk->ik', np.einsum('nij,jk->nik', A.transpose(0,2,1), B), A)
100 loops, best of 3: 10.7 ms per loop

In [132]: %timeit np.einsum('nij, il, kln ->jk', A, B, A.T)
10 loops, best of 3: 105 ms per loop

In [133]: %timeit np.tensordot(A, np.tensordot(A, B, axes=(1, 1)), axes=((0, 1), (0, 2)))
100 loops, best of 3: 6.22 ms per loop

In [134]: %timeit ((A.transpose(2,0,1).reshape(-1,m).dot(B.T)).reshape(n,-1)).dot(A.reshape(-1,n)).T
100 loops, best of 3: 5.3 ms per loop
于 2015-10-28T22:28:45.003 回答