Question

我想知道如何转换此问题以减少代码中np.sum()函数调用的开销。

我有一个input矩阵，比如说shape=(1000, 36)。每行代表图中的一个节点。我有一个我正在做的操作，它迭代每一行并对可变数量的其他行进行元素添加。这些“其他”行在字典nodes_nbrs中定义，为每行记录必须汇总在一起的行列表。一个例子就是这样：

nodes_nbrs = {0: [0, 1], 
              1: [1, 0, 2],
              2: [2, 1],
              ...}

此处，节点0将转换为节点0和1的总和。节点1将转换为节点1，0和2的总和。等等其他节点。

我目前实施的当前（和天真）方式就是这样。我首先实例化我想要的最终形状的零数组，然后迭代nodes_nbrs字典中的每个键值对：

output = np.zeros(shape=input.shape)
for k, v in nodes_nbrs.items():
    output[k] = np.sum(input[v], axis=0)

这个代码在小测试（shape=(1000, 36)）中都很酷很好，但在较大的测试（shape=(~1E(5-6), 36)）上，完成需要2-3秒。我最终不得不做几千次这样的操作，所以我试着看看是否有更优化的方法来做到这一点。

在进行线性剖析后，我注意到这里的关键杀手是反复调用np.sum函数，这占用了总时间的50％左右。有没有办法可以消除这种开销？还是有另外一种方法可以优化它吗？

除此之外，这里列出了我已经完成的事情，并且（非常简要地说）了他们的结果：

cython版本：消除了for循环类型检查开销，减少了30％的时间。使用cython版本时，np.sum占整个挂钟时间的80％左右，而不是50％。
将np.sum预先声明为变量npsum，然后在npsum循环内调用for。与原版没什么区别。
将np.sum替换为np.add.reduce，并将其分配给变量npsum，然后在npsum循环内调用for。挂钟时间缩短约10％，但与autograd不兼容（稀疏矩阵子弹点下面的解释）。
numba JIT-ing：没有尝试过添加装饰器。没有改善，但没有更努力。
将nodes_nbrs字典转换为密集的numpy二进制数组（1和0），然后执行单np.dot次操作。理论上很好，在实践中很糟糕，因为它需要一个shape=(10^n, 10^n)的方阵，这在内存使用方面是二次的。

我没有尝试过的事情，但我却犹豫不决：

scipy稀疏矩阵：我使用的是autograd，它不支持自动区分dot稀疏矩阵的scipy操作。

对于那些好奇的人来说，这实际上是对图结构数据的卷积运算。有点有趣的是为毕业学校开发这个，但也有点令人沮丧的知识的最前沿。

Answer 1

如果scipy.sparse不是一个选项，你可能采用的一种方法是按摩数据，这样你就可以使用矢量化函数来完成编译层中的所有操作。如果将邻居字典更改为具有缺失值的适当标志的二维数组，则可以使用np.take提取所需数据，然后执行单个sum()调用。

以下是我想到的一个例子：

import numpy as np

def make_data(N=100):
    X = np.random.randint(1, 20, (N, 36))
    connections = np.random.randint(2, 5, N)
    nbrs = {i: list(np.random.choice(N, c))
            for i, c in enumerate(connections)}
    return X, nbrs

def original_solution(X, nbrs):
    output = np.zeros(shape=X.shape)
    for k, v in nbrs.items():
        output[k] = np.sum(X[v], axis=0)
    return output

def vectorized_solution(X, nbrs):
    # Make neighbors all the same length, filling with -1
    new_nbrs = np.full((X.shape[0], max(map(len, nbrs.values()))), -1, dtype=int)
    for i, v in nbrs.items():
        new_nbrs[i, :len(v)] = v

    # add a row of zeros to X
    new_X = np.vstack([X, 0 * X[0]])

    # compute the sums
    return new_X.take(new_nbrs, 0).sum(1)

现在我们可以确认结果是否匹配：

>>> X, nbrs = make_data(100)
>>> np.allclose(original_solution(X, nbrs),
                vectorized_solution(X, nbrs))
True

我们可以花时间看看加速：

X, nbrs = make_data(1000)
%timeit original_solution(X, nbrs)
%timeit vectorized_solution(X, nbrs)
# 100 loops, best of 3: 13.7 ms per loop
# 100 loops, best of 3: 1.89 ms per loop

升级到更大尺寸：

X, nbrs = make_data(100000)
%timeit original_solution(X, nbrs)
%timeit vectorized_solution(X, nbrs)
1 loop, best of 3: 1.42 s per loop
1 loop, best of 3: 249 ms per loop

它的速度提高了5到10倍，这对于您的目的来说可能已经足够了（尽管这很大程度上取决于您的nbrs词典的确切特征）。

修改：为了好玩，我尝试了其他一些方法，一个使用numpy.add.reduceat，一个使用pandas.groupby，另一个使用scipy.sparse。似乎我最初提出的矢量化方法可能是最好的选择。这里他们是供参考：

from itertools import chain

def reduceat_solution(X, nbrs):
    ind, j = np.transpose([[i, len(v)] for i, v in nbrs.items()])
    i = list(chain(*(nbrs[i] for i in ind)))
    j = np.concatenate([[0], np.cumsum(j)[:-1]])
    return np.add.reduceat(X[i], j)[ind]

np.allclose(original_solution(X, nbrs),
            reduceat_solution(X, nbrs))
# True

-

import pandas as pd

def groupby_solution(X, nbrs):
    i, j = np.transpose([[k, vi] for k, v in nbrs.items() for vi in v])
    return pd.groupby(pd.DataFrame(X[j]), i).sum().values

np.allclose(original_solution(X, nbrs),
            groupby_solution(X, nbrs))
# True

-

from scipy.sparse import csr_matrix
from itertools import chain

def sparse_solution(X, nbrs):
    items = (([i]*len(col), col, [1]*len(col)) for i, col in nbrs.items())
    rows, cols, data = (np.array(list(chain(*a))) for a in zip(*items))
    M = csr_matrix((data, (rows, cols)))
    return M.dot(X)

np.allclose(original_solution(X, nbrs),
            sparse_solution(X, nbrs))
# True

所有的时间安排在一起：

X, nbrs = make_data(100000)
%timeit original_solution(X, nbrs)
%timeit vectorized_solution(X, nbrs)
%timeit reduceat_solution(X, nbrs)
%timeit groupby_solution(X, nbrs)
%timeit sparse_solution(X, nbrs)
# 1 loop, best of 3: 1.46 s per loop
# 1 loop, best of 3: 268 ms per loop
# 1 loop, best of 3: 416 ms per loop
# 1 loop, best of 3: 657 ms per loop
# 1 loop, best of 3: 282 ms per loop

Answer 2

基于最近稀疏问题的工作，例如Extremely slow sum row operation in Sparse LIL matrix in Python

这里是如何用稀疏矩阵解决你的问题。该方法可能同样适用于密集的方法。这个想法是稀疏sum实现为矩阵乘积，行（或列）为1。稀疏矩阵的索引很慢，但矩阵乘积是很好的C代码。

在这种情况下，我将构建一个乘法矩阵，对于我想要求和的行，它有1个 - 字典中每个条目的1个不同的集合。

样本矩阵：

In [302]: A=np.arange(8*3).reshape(8,3)    
In [303]: M=sparse.csr_matrix(A)

选择词典：

In [304]: dict={0:[0,1],1:[1,0,2],2:[2,1],3:[3,4,7]}

从这本词典中构建一个稀疏矩阵。这可能不是构建这样一个矩阵的最有效方法，但它足以证明这个想法。

In [305]: r,c,d=[],[],[]
In [306]: for i,col in dict.items():
    c.extend(col)
    r.extend([i]*len(col))
    d.extend([1]*len(col))

In [307]: r,c,d
Out[307]: 
([0, 0, 1, 1, 1, 2, 2, 3, 3, 3],
 [0, 1, 1, 0, 2, 2, 1, 3, 4, 7],
 [1, 1, 1, 1, 1, 1, 1, 1, 1, 1])

In [308]: idx=sparse.csr_matrix((d,(r,c)),shape=(len(dict),M.shape[0]))

执行求和并查看结果（作为密集数组）：

In [310]: (idx*M).A
Out[310]: 
array([[ 3,  5,  7],
       [ 9, 12, 15],
       [ 9, 11, 13],
       [42, 45, 48]], dtype=int32)

这是原始的比较。

In [312]: M.A
Out[312]: 
array([[ 0,  1,  2],
       [ 3,  4,  5],
       [ 6,  7,  8],
       [ 9, 10, 11],
       [12, 13, 14],
       [15, 16, 17],
       [18, 19, 20],
       [21, 22, 23]], dtype=int32)

代码优化 - Python中的函数调用数

2 个答案: