哈,只需一行即可完成:np.einsum('nmk,nkj->mj',A,B)。在
不要将变量命名为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
尺寸越大,事情就越不一样了:
^{pr2}$
以及: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