python 中稀疏矩阵的并行计算

Parallel calculation on a sparse matrix in python

我有一个很大的 [numpy] 向量 X 和一个比较函数 f(x,y)。我需要找到 X 的所有元素对 f(X[I],X[j])<T 对于某个阈值 T。这很有效:

good_inds = {}
for i in range(0,len(X)):
   for j in range(x+1,len(X)):
       score = f(X[i],X[j])
       if score<T:
           good_inds[x,y] = score

这实际上构建了一个表示稀疏矩阵的字典。问题是它相当慢,我希望将这个过程并行化。 请指教

这很适合 multiprocessing.Pool

创建你的 numpy 数组,然后创建一个包含所有可能 ij 值的二元组的迭代器。例如 itertools.combinations.

In [1]: import itertools

In [7]: list(itertools.combinations(range(4), 2))                                                        
Out[7]: [(0, 1), (0, 2), (0, 3), (1, 2), (1, 3), (2, 3)]

(当然,您应该使用向量的长度作为 range 的参数。)

编写如下函数:

def worker(pair):
    i, j = pair
    rv = False
    if f(X[i],X[j]) < T:
       rv = True
    return (i, j, rv)

创建一个Pool,然后运行 imap_unordered:

p = multiprocessing.Pool()
for i, j, result in p.imap_unordered(worker, itertools.combinations(range(len(X)), 2)):
    if result:
        print('Good pair:', i, j)
        # do something with the results...

这将 运行 与你 CPU 拥有的核心数一样多的工人。

所以。显然 SciPy 已经足够好了。

full_dist_mat = spatial.distance.squareform( spatial.distance.pdist(vects2, metric='cosine'))

已经优化。 运行 在 Macbook pro 上的 jupyter lab 中,2000 个向量需要 1.3 秒。