Question

我希望使用Numba（http://numba.pydata.org/）来加速我的机器学习算法（用Python编写）。注意，该算法将稀疏矩阵作为其输入数据。在我的纯Python实现中，我使用了Scipy的csr_matrix和相关类，但显然它与Numba的JIT编译器不兼容。

我还创建了自己的自定义类来实现稀疏矩阵（基本上是（索引，值）对列表的列表），但它再次与Numba不兼容（即，我得到一些奇怪的错误消息说它无法识别扩展类型）

是否有另一种简单的方法来实现稀疏矩阵，只使用与Numba兼容的numpy（不使用SciPy）？任何示例代码将不胜感激。谢谢！

Answer 1

您可以将稀疏矩阵的数据作为纯numpy或python访问。例如

M=sparse.csr_matrix([[1,0,0],[1,0,1],[1,1,1]])
ML = M.tolil()

for d,r in enumerate(zip(ML.data,ML.rows))
    # d,r are lists
    dr = np.array([d,r])
    print dr

产生

[[1]
 [0]]
[[1 1]
 [0 2]]
[[1 1 1]
 [0 1 2]]

当然，numba可以处理使用这些数组的代码，当然，前提是它不希望每行具有相同大小的数组。

Answer 2

如果您要做的只是遍历CSR矩阵的值，则可以将属性数据，indptr和索引传递给函数而不是CSR矩阵对象。

from scipy import sparse
from numba import njit

@njit
def print_csr(A, iA, jA):
    for row in range(len(iA)-1):
        for i in range(iA[row], iA[row+1]):
            print(row, jA[i], A[i])

A = sparse.csr_matrix([[1, 2, 0], [0, 0, 3], [4, 0, 5]])
print_csr(A.data, A.indptr, A.indices)

Numba中的稀疏矩阵

2 个答案: