稀疏矩阵索引操作的向量化

Question

以下代码运行得太慢了，尽管看起来所有操作都已经进行了向量化处理。

from numpy import *
from scipy.sparse import *

n = 100000;
i = xrange(n); j = xrange(n);
data = ones(n);

A=csr_matrix((data,(i,j)));

x = A[i,j]

问题似乎出在索引操作是用Python函数实现的，当你调用 A[i,j] 时，会产生以下性能分析的结果。

         500033 function calls in 8.718 CPU seconds

   Ordered by: internal time

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
   100000    7.933    0.000    8.156    0.000 csr.py:265(_get_single_element)
        1    0.271    0.271    8.705    8.705 csr.py:177(__getitem__)
(...)

也就是说，Python函数 _get_single_element 被调用了100000次，这样效率真的很低。为什么不直接用纯C语言来实现呢？有没有人知道有什么方法可以绕过这个限制，加快上面的代码运行速度？我是否应该使用不同的稀疏矩阵类型？

Answer 1

我觉得只有在使用默认的数据类型'object'时，_get_single_element这个函数才会被调用。你有没有试过指定一个数据类型，比如用csr_matrix((data, (i,j)), dtype=int32)？

Answer 2

你可以用 A.diagonal() 来更快地获取对角线元素（在我的电脑上，速度是0.0009秒，而不是3.8秒）。不过，如果你想随机访问一些元素，那就比较复杂了，因为这不是简单的切片，而是需要用到一系列的索引。_get_single_element 这个函数被调用了100000次，因为它在遍历你传给它的迭代器（i和j）。切片的例子是 A[30:60,10] 或类似的东西。

另外，我建议用 csr_matrix(eye(n,n)) 来创建你用迭代器做的同样的矩阵，这样会简单一些。

更新：

好的，既然你的问题是关于如何快速访问很多随机元素，我会尽量回答你的问题。

• 为什么这个没有用纯C实现？

答案很简单：还没有人去做这个。从我看到的情况来看，Scipy的稀疏矩阵模块还有很多工作要做。其中一部分已经用C实现了，就是不同稀疏矩阵格式之间的转换。

• 有没有人知道如何绕过这个限制，加速上面的代码？

你可以尝试深入稀疏矩阵模块，看看能不能加速。我试过，结果把你上面代码的随机访问时间缩短到了原来的三分之一。我直接访问了 _get_single_element，并大幅简化了代码，包括去掉了一些边界检查。

不过，使用 lil_matrix 还是更快（虽然初始化矩阵的速度较慢），但我必须用列表推导式来访问，因为 lil_matrix 不支持你那种索引方式。使用列表推导式访问 csr_matrix 的速度还是比 lil_matrix 慢。总的来说，lil_matrix 在随机访问元素时更快，因为它没有压缩。

使用原始形式的 lil_matrix 运行时间大约是你上面代码的五分之一。如果我去掉一些边界检查，并直接调用 lil_matrix 的 _get1() 方法，时间还能再缩短大约7%，也就是从3.4-3.8秒降到大约0.261秒。

最后，我尝试自己写一个函数，直接访问 lil_matrix 的数据，避免重复调用函数。这个方法的时间大约是0.136秒。这没有利用数据的排序，这也是一个潜在的优化（特别是当你访问很多在同一行的元素时）。

如果你想要更快的速度，那可能需要自己写一个C语言的稀疏矩阵实现。

• 我应该使用不同类型的稀疏矩阵吗？

如果你打算访问很多元素，我建议使用 lil_matrix，但这也要看你具体需要做什么。比如，你是否还需要进行矩阵乘法？记住，矩阵之间的转换在某些情况下可能会很快，所以不要排除为了不同的操作而换成其他矩阵格式。

如果你不需要对矩阵进行任何代数运算，那也许可以考虑使用 defaultdict 或类似的东西。使用 defaultdict 的风险在于，当请求一个不在字典中的元素时，它会将该项设置为默认值并存储，这可能会导致问题。