如何在Python中加速矩阵乘法？

Question

我现在正在编写一个模拟程序，并且在优化速度。通过查看性能分析工具line_profiler，我发现我的瓶颈在于矩阵乘法（np.matmul）。我需要对一个矩阵'A'（大小为[990, 1, 10, 3, 3]）和一个矩阵'B'（大小为[990, 1, 10, 3, 1]）进行大约80万次不同的乘法运算，而每次乘法的A和B都是不同的。下面是一些简化的“测试”代码（在这个简单的例子中，矩阵保持不变，所以请不要太关注这一点）：

precession = np.random.uniform(1,10,[990, 1, 10, 3, 3])
vecMblood = np.random.uniform(1,10,[990, 1, 10, 3, 1])

start = timeit.default_timer()
for _ in range(4000*200):
    np.matmul(precession, vecMblood)
end = timeit.default_timer()
print(f"np.matmul took {end - start} seconds.")

在单核处理器上，这个过程需要133秒。我在主代码中使用了多线程，这确实提高了一些速度——也就是说，在原始代码中，有4000次乘法是顺序进行的，但在200个测试案例中是独立的——因此这4000次乘法在12个核心上顺序运行，总共是4000x200（800,000）。我还将主代码中的数据类型从float64减少到float32，这也有所帮助。不过，除了np.matmul之外，Python中有没有更优化的方式来处理这个问题？我也尝试在其他地方实现Numba函数，但开销太大了——所以我在寻找一种没有这种开销的方法。

Answer 1

你可以通过使用 einsum 来让这个过程快大约40%。

np.einsum("...ij,...j", precession, vecMblood[..., 0])[..., None]

时间记录：

import numpy as np

rng = np.random.default_rng()

a = rng.random([990, 1, 10, 3, 3])
b = rng.random(([990, 1, 10, 3, 1]))

%timeit -n 1000 np.matmul(a, b)
%timeit -n 1000 np.einsum("...ij,...j", a, b[..., 0])[..., None]

输出结果：

190 µs ± 10 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
108 µs ± 816 ns per loop (mean ± std. dev. of 7 runs, 1,000 loops each)