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

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(稀疏矩阵子弹点下面的解释).
• numbaJIT-ing:没有尝试过添加装饰器.没有改善,但没有更努力.
• 将nodes_nbrs字典转换为密集numpy二进制数组(1s和0s),然后执行单个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