如何将产生稀疏矩阵的代码矢量化？

Question

Let me start off by explaining what I want to do. I am trying to build a recommendation system based off of m packages, each with n features, stored in an mxn sparse matrix X . To do this, I'm attempting to run kNN to get the k closest matches for a packages. I want to build an mxm sparse matrix K where K[i, j] is the dot product of rows X[i] and X[j] if X[j] was a package returned by kNN for X[i] , otherwise 0.

Here is the code I've written:

X = ...
knn = NearestNeighbors(n_neighbors=self.n_neighbors, metric='l2')
knn.fit(X)
knn_indices = knn.kneighbors(X, return_distance=False)

m, k = X.shape[0], self.n_neighbors
K = lil_matrix((m, m))

for i, indices in enumerate(knn_indices):
    xi = X.getrow(i)
    for j in indices:
        xj = X.getrow(j)
        K[i, j] = xi.dot(xj.T)[0, 0]

I'm trying to figure out how to make this more efficient. In my scenario, m is ~1.2 million, n is ~50000, and k is 500, so perf is very important.

The last part where I populate K is the bottleneck of my program. getrow seems to perform very poorly; according to the scipy docs, it makes a copy of the row, so getrow call could be copying up to 50k elements each time it's called. Also, in the innermost loop I can't figure out how to get back a scalar for dot instead of creating a whole new 1x1 sparse matrix.

How can I avoid these problems and speed up/vectorize the last part of this code? Thanks.

Answer 1

In [21]: from scipy import sparse
In [22]: M = sparse.random(10,10,.2,'csr')
In [23]: M
Out[23]: 
<10x10 sparse matrix of type '<class 'numpy.float64'>'
    with 20 stored elements in Compressed Sparse Row format>

Looking a MA , I selected this small knn_indices array for testing:

In [45]: knn = np.array([[4],[2],[],[1,3]])

Your double loop:

In [46]: for i, indices in enumerate(knn):
    ...:     xi = M[i,:]
    ...:     for j in indices:
    ...:         xj = M[j,:]
    ...:         print((xi*xj.T).A)
    ...:         
[[0.35494592]]
[[0.]]
[[0.08112133]]
[[0.56905781]]

The inner loop can be condensed:

In [47]: for i, indices in enumerate(knn):
    ...:     xi = M[i,:]
    ...:     xj = M[indices,:]
    ...:     print((xi*xj.T).A)
    ...:         
[[0.35494592]]
[[0.]]
[]
[[0.08112133 0.56905781]]

and with the assignment:

In [49]: k = sparse.lil_matrix((4,5))
In [50]: for i, indices in enumerate(knn):
    ...:     xi = M[i,:]
    ...:     for j in indices:
    ...:         xj = M[j,:]
    ...:         k[i,j] = (xi*xj.T)[0,0]
    ...:         
    ...:         
In [51]: k.A
Out[51]: 
array([[0.        , 0.        , 0.        , 0.        , 0.35494592],
       [0.        , 0.        , 0.        , 0.        , 0.        ],
       [0.        , 0.        , 0.        , 0.        , 0.        ],
       [0.        , 0.08112133, 0.        , 0.56905781, 0.        ]])

The second loop with

 k[i,indices] = (xi*xj.T)

does the same thing.

It may be possible to do something with the i loop as well, but this is at least a start.

That knn doesn't need to an array. With differing inner list lengths it's an object dtype anyways. Better leave it as list.

An alternative to filling this lil matrix, would be to accumulate i , indices and the dot product in coo style arrays.

In [64]: r,c,d = [],[],[]
In [65]: for i, indices in enumerate(knn):
    ...:     xi = M[i,:]
    ...:     xj = M[indices,:]
    ...:     t = (xi*xj.T).data
    ...:     if len(t)>0:
    ...:         r.extend([i]*len(indices))
    ...:         c.extend(indices)
    ...:         d.extend(t)
    ...:         
In [66]: r,c,d
Out[66]: 
([0, 3, 3],
 [4, 1, 3],
 [0.3549459176547072, 0.08112132851228658, 0.5690578146292733])
In [67]: sparse.coo_matrix((d,(r,c))).A
Out[67]: 
array([[0.        , 0.        , 0.        , 0.        , 0.35494592],
       [0.        , 0.        , 0.        , 0.        , 0.        ],
       [0.        , 0.        , 0.        , 0.        , 0.        ],
       [0.        , 0.08112133, 0.        , 0.56905781, 0.        ]])

In my test case the 2nd row doesn't have any nonzero values, requiring an extra test in the loop. I don't know if this is any faster than the lil approach.