如何在 Python 中更快地进行矩阵乘法？ [重复]

Question

我目前正在编写模拟并优化速度。看看 line_profiler，我的瓶颈在于矩阵乘法（np.matmul）。我正在对矩阵 'A'([990, 1, 10, 3, 3]) 和矩阵 'B' ([990, 1, 10, 3, 1]) 执行约 800,000 次不同的乘法，其中 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 次乘法在 200 个用例的 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)