在numpy中将小矩阵与标量相乘的最有效方法

Question

我有一个程序，其主要性能瓶颈涉及将一个维度为 1 的矩阵和另一个大维度的矩阵相乘，例如1000：

large_dimension = 1000

a = np.random.random((1,))
b = np.random.random((1, large_dimension))

c = np.matmul(a, b)

换句话说，将矩阵 b 与标量 a[0] 相乘。

我正在寻找最有效的计算方法，因为这个操作重复了数百万次。

我测试了这两种简单方法的性能，它们实际上是等效的：

%timeit np.matmul(a, b)
>> 1.55 µs ± 45.8 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

%timeit a[0] * b
>> 1.77 µs ± 34.6 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

有没有更有效的计算方法？

• 注意：我无法将这些计算转移到 GPU，因为程序正在使用多处理并且许多此类计算是并行完成的。

Answer 1

在这种情况下，使用逐元素乘法可能会更快，但您看到的时间主要是 Numpy 的开销（从 CPython 解释器调用 C 函数，包装/展开类型，进行检查，执行操作，数组分配等）。

因为这个操作重复了数百万次

这就是问题所在。事实上，CPython 解释器 非常不擅长以低延迟 处理事情。当您处理 Numpy 类型时尤其如此，因为调用 C 代码并执行检查琐碎操作比在纯 Python 中执行它要慢得多，这也比编译的本机 C/C++ 代码慢得多。如果你真的需要这个，并且你不能使用 Numpy矢量化你的代码（因为你有一个循环迭代时间步长），那么你就不要使用 CPython，或者至少不要使用纯 Python 代码。相反，您可以使用 Numba 或 Cython 来减轻执行 C 调用、包装类型等的影响。如果这还不够，那么您将需要编写原生 C /C++ 代码（或任何类似的语言），除非你找到一个专用的 Python 包完全为你做这件事。请注意，Numba 仅在适用于本机类型或 Numpy 数组（包含本机类型）时才快速。如果您使用大量纯 Python 类型并且不想重写代码，那么您可以尝试 PyPy JIT。

---

这是 Numba 中的一个简单示例，可避免（昂贵的）创建/分配新数组（以及许多 Numpy 内部检查和调用），专门用于解决您的特定情况：

@nb.njit('void(float64[::1],float64[:,::1],float64[:,::1])')
def fastMul(a, b, out):
    val = a[0]
    for i in range(b.shape[1]):
        out[0,i] = b[0,i] * val

res = np.empty(b.shape, dtype=b.dtype)
%timeit fastMul(a, b, res)
# 397 ns ± 0.587 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

在撰写本文时，此解决方案比其他所有解决方案都快。由于大部分时间都花在调用 Numba 和执行一些内部检查上，因此将 Numba 直接用于包含迭代循环的函数应该会产生更快的代码。

Answer 2

large_dimension = 1000

a = np.random.random((1,))
B = np.random.random((1, large_dimension))

%timeit np.matmul(a, B)
5.43 µs ± 22 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

%timeit a[0] * B
5.11 µs ± 6.92 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

只使用浮动

%timeit float(a[0]) * B
3.48 µs ± 26.1 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

为了避免内存分配使用“缓冲区”

buffer = np.empty_like(B)

%timeit np.multiply(float(a[0]), B, buffer)
2.96 µs ± 37.1 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

为避免不必要的获取属性使用“别名”

mul = np.multiply

%timeit mul(float(a[0]), B, buffer)
2.73 µs ± 12.6 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

而且我完全不建议使用 numpy 标量， 因为如果你避免它，计算会更快

a_float = float(a[0])

%timeit mul(a_float, B, buffer)
1.94 µs ± 5.74 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

此外，如果有可能，则在循环外初始化缓冲区一次（当然，如果你有类似循环的东西：）

rng = range(1000)

%%timeit
for i in rng:
    pass
24.4 µs ± 1.21 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

%%timeit
for i in rng:
    mul(a_float, B, buffer)
1.91 ms ± 2.21 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

所以，

"best_iteration_time" = (1.91 - 0.02) / 1000 => 1.89 (µs)

“加速” = 5.43 / 1.89 = 2.87

Answer 3

import numpy as np
import numba

def matmult_numpy(matrix, c):
    return np.matmul(c, matrix)

@numba.jit(nopython=True)
def matmult_numba(matrix, c):
    return c*matrix

if __name__ == "__main__":
    large_dimension = 1000
    a = np.random.random((1, large_dimension))
    c = np.random.random((1,))

使用 Numba 的速度大约提高了 3 倍。 Numba cognoscenti 可以通过将参数“c”显式转换为标量来做得更好

检查：结果

%timeit matmult_numpy(a, c) 2.32 µs 每循环 ± 50 ns（平均值 ± 标准偏差，7 次运行，每次 100000 次循环）

%timeit matmult_numba(a, c) 763 ns 每个循环 ± 6.67 ns（平均值 ± 标准偏差，7 次运行，每次 1000000 次循环）