在 Numpy 中“矢量化”大量矩阵运算的最快方法是什么？答案

【问题标题】：What is the fastest way to "vectorize" a large number of matrix operations in Numpy?在 Numpy 中“矢量化”大量矩阵运算的最快方法是什么？
【发布时间】：2021-09-09 21:31:41
【问题描述】：

假设我想在 Numpy 中做很多矩阵乘法；最快的方法是什么？

具体来说，这就是问题所在：我有两个很长的矩阵列表，我想将它们按元素相乘。也就是说，我有

[a_1, a_2, a_3, ..., a_N]

和

[b_1, b_2, b_3, ..., b_N],

每个a_i、b_i 是一个nxn 矩阵（n 很小，比如n=2），N 很大（比如N = 100000），我想找到矩阵积a_1 * b_1, a_2 * b_2, ...

使用 Python 和 Numpy/Scipy 最快的方法是什么？

一些选项是：

带有for 循环——这很慢，因为 Python 循环很慢。
将小矩阵放入两个NxN 块对角矩阵A 和B--这将导致必须乘以比所需更大的矩阵。
使用vectorize-- 这是最容易编码的，但并不比for 循环快。

【问题讨论】：

你看过np.vectorize吗？
@PranavVempati vectorize 不会加快速度，它与自己编写 for 循环的速度相同。
@PranavVempati np.vectorize 基本上是一个 Python 循环
"vectorize" 在 numpy 加速意义上实际上意味着用编译代码中的循环替换 python 级别循环，通常使用现有的 numpy 方法。答案使用matmul的“批处理”能力。

标签： python numpy linear-algebra

【解决方案1】：

您已经可以乘以 3D 数组，只需将您的数组列表放入 numpy 数组中，例如，

A = np.array([a_1, a_2, ..., a_N])
B = np.array([b_1, b_2, ..., b_N])

然后乘以A @ B（@ 是矩阵乘法运算符）。这是一个使用两个 3x3 数组“列表”的示例：

In [1]: import numpy as np

In [2]: x = np.random.randint(0, 9, (2, 3, 3))

In [3]: y = np.random.randint(0, 9, (2, 3, 3))

In [4]: x
Out[4]:
array([[[0, 4, 8],
        [2, 5, 5],
        [3, 0, 5]],

       [[7, 6, 1],
        [7, 0, 7],
        [5, 2, 8]]])

In [5]: y
Out[5]:
array([[[7, 2, 6],
        [6, 1, 4],
        [6, 8, 5]],

       [[8, 5, 4],
        [8, 2, 7],
        [3, 7, 0]]])

In [7]: x @ y
Out[7]:
array([[[ 72,  68,  56],
        [ 74,  49,  57],
        [ 51,  46,  43]],

       [[107,  54,  70],
        [ 77,  84,  28],
        [ 80,  85,  34]]])

为了证明这一切都是每个矩阵在相应索引处的乘积：

In [8]: x[0]
Out[8]:
array([[0, 4, 8],
       [2, 5, 5],
       [3, 0, 5]])

In [9]: y[0]
Out[9]:
array([[7, 2, 6],
       [6, 1, 4],
       [6, 8, 5]])

In [10]: x[0] @ y[0]
Out[10]:
array([[72, 68, 56],
       [74, 49, 57],
       [51, 46, 43]])

In [11]: (x @ y)[0]
Out[11]:
array([[72, 68, 56],
       [74, 49, 57],
       [51, 46, 43]])

【讨论】：

【解决方案2】：

只需像往常一样使用numpy.matmul 或@，它是一个ufunc，可以进行广播，但不是在数组元素上，而是在矩阵子数组上。在您的情况下，您只需将 (n,m) 矩阵堆叠在 (N,n,m) numpy 数组中，将 (m,p) 矩阵堆叠在 (N,m,p) numpy 数组中。

m = np.array([[1, 3], [2,4]])

m
Out[12]: 
array([[1, 3],
       [2, 4]])

m @ m
Out[13]: 
array([[ 7, 15],
       [10, 22]])

stackedm = np.stack([m,m,m])

stackedm
Out[15]: 
array([[[1, 3],
        [2, 4]],

       [[1, 3],
        [2, 4]],

       [[1, 3],
        [2, 4]]])

stackedm @ stackedm
Out[16]: 
array([[[ 7, 15],
        [10, 22]],

       [[ 7, 15],
        [10, 22]],

       [[ 7, 15],
        [10, 22]]])

【讨论】：