在numpy中获得矩阵<矩阵结果的最快方法是什么？

Question

假设我有一个维度为（M，A）的矩阵M_1和一个维度为（M，B）的矩阵M_2。 M_1 < M_2 的结果应该是一个维度为（M，B，A）的矩阵，其中每一行M1都与相应行的M_2中的每个元素进行比较，并为每个比较给出一个布尔向量（或1,0向量）。

例如，如果我有一个矩阵：

M1 = [[1,2,3]
      [3,4,5]]

M2 = [[1,2],
      [3,4]]

result should be [[[False, False, False],
                   [True, False, False]],
                  [[False, False, False], 
                   [True, False, False]]]

目前，我正在使用for循环，当我需要重复这些操作多次（需要几个月）时，速度非常慢。希望有一种向量化的方法来完成此操作。如果没有，我还能做什么？

我正在查看M_1为（500，3000000），M_2为（500，500），并且需要重复大约10000次。

- user10024395

---

@Divakar 抱歉，已编辑。 - user10024395

由于这是一个完全受内存限制的问题（如@Divakar所示的numexpr，或者如果使用Numba或Cython），所以了解您接下来将如何处理此布尔数组会很有帮助。 - max9111

我不相信你真的需要这个。 - Veedrac

一个 (500, 500, 3000000) 的布尔矩阵大小接近 700 GiB。采用位表示法可以将其减小到不到 90 GiB，但这意味着您不能使用基本的 Numpy 矩阵。 - Veedrac

那么矩阵代表什么，M和B又是什么？（我知道A是图像数量。）这似乎与通常的做法完全不同，通常的做法是仅计算准确性；请参阅标准论文。 - Veedrac

Answer 1

针对NumPy数组，使用None/np.newaxis扩展dims以使第一个轴对齐，而第二个轴是扩展的，这样可以逐元素比较它们。最后通过broadcasting来进行比较，以实现向量化的解决方案。

M1[:,None,:] < M2[:,:,None]

样例运行 -

In [19]: M1
Out[19]: 
array([[1, 2, 3],
       [3, 4, 5]])

In [20]: M2
Out[20]: 
array([[1, 2],
       [3, 4]])

In [21]: M1[:,None,:] < M2[:,:,None]
Out[21]: 
array([[[False, False, False],
        [ True, False, False]],

       [[False, False, False],
        [ True, False, False]]])

对于列表作为输入，请使用numpy.expand_dims，然后进行比较 -

In [42]: M1 = [[1,2,3],
    ...:       [3,4,5]]
    ...: 
    ...: M2 = [[1,2],
    ...:       [3,4]]

In [43]: np.expand_dims(M1, axis=1) < np.expand_dims(M2, axis=2)
Out[43]: 
array([[[False, False, False],
        [ True, False, False]],

       [[False, False, False],
        [ True, False, False]]])

进一步提升

借助numexpr模块的多核支持，可以进一步提升处理大型数据时的性能 -

In [44]: import numexpr as ne

In [52]: M1 = np.random.randint(0,9,(500, 30000))

In [53]: M2 = np.random.randint(0,9,(500, 500))

In [55]: %timeit M1[:,None,:] < M2[:,:,None]
1 loop, best of 3: 3.32 s per loop

In [56]: %timeit ne.evaluate('M1e<M2e',{'M1e':M1[:,None,:],'M2e':M2[:,:,None]})
1 loop, best of 3: 1.53 s per loop

- Divakar