如何加速numpy中的乘法和求和运算 [英] How to speed up multiply and sum operations in numpy

问题描述

我需要解决一个有限元方法问题，并且必须从 A 和 B 计算以下 C 和一个大的 M (M>1M).例如，

I need to solve a Finite Element Method problem and have to calculate the following C from A and B with a large M (M>1M). For example,

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 和 C 的轴可以自由变形或重新排列.

Could anyone come up with a code which is faster than (A * B).sum(axis = -1)? You can reshape or re-arrange the axes of A, B, and C freely.

推荐答案

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

You can use np.einsum for a slightly more efficient approach, both in performance and memory usage:

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

这篇关于如何加速numpy中的乘法和求和运算的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！