如何加速 numpy 中的乘法和求和运算[重复]

我需要解决有限元法问题并且必须计算以下内容C from A and B与一个大M (M>1M）。例如，

import numpy as np
M=4000000
A=np.random.rand(4, M, 3)
B=np.random.rand(M,3)
C = (A * B).sum(axis = -1) # need to be optimized

任何人都可以想出一个比(A * B).sum(axis = -1)？您可以重塑或重新排列轴A, B, and C freely.

您可以使用np.einsum对于性能和内存使用方面稍微更有效的方法：

M=40000
A=np.random.rand(4, M, 3)
B=np.random.rand(M,3)
out = (A * B).sum(axis = -1) # need to be optimized

%timeit (A * B).sum(axis = -1) # need to be optimized
# 5.23 ms ± 198 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

%timeit np.einsum('ijk,jk->ij', A, B)
# 1.31 ms ± 136 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

np.allclose(out, np.einsum('ijk,jk->ij', A, B))
# True