代碼優化-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調用np.sum函數，這大約占總時間的50％。 有沒有辦法可以消除這種開銷？ 還是我可以優化它的另一種方法？

---

除此之外，這是我所做的事情的清單，以及它們的結果（非常簡短）：

• cython版本：消除了for循環類型檢查的開銷，減少了30％的時間。 使用cython版本時， np.sum大約占總掛鍾時間的80％，而不是50％。
• 將np.sum預先聲明為變量npsum ，然后在for循環內調用npsum 。 與原件沒有區別。
• 將np.sum替換為np.add.reduce ，並將其分配給變量npsum ，然后在for循環內調用npsum 。 掛鍾時間減少了約10％，但與autograd不兼容（下面在稀疏矩陣項目符號中進行了說明）。
• numba JIT-ing：除了添加裝飾器，沒有其他嘗試。 沒有改善，但是沒有努力。
• 將nodes_nbrs詞典轉換為密集的numpy二進制數組（1s和0s），然后執行單個np.dot操作。 理論上不錯，實踐上很不好，因為它需要一個shape=(10^n, 10^n)方陣，這在內存使用方面是二次方的。

我沒有嘗試過的事情，但猶豫不決：

• scipy稀疏矩陣：我正在使用autograd ，它不支持對scipy稀疏矩陣進行dot運算的自動區分。

---

對於那些好奇的人，這本質上是對圖結構數據的卷積運算。 Kinda很高興為研究生學校開發此軟件，但在知識的最前沿也有些沮喪。

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

基於最近稀疏問題的工作，例如，Python的稀疏LIL矩陣中的極慢和行操作

這是使用稀疏矩陣可以解決您的問題的方式。 該方法可能也適用於密集型方法。 這個想法是將稀疏sum實現為行（或列）為1s的矩陣乘積。 稀疏矩陣的索引編制很慢，但矩陣乘積是良好的C代碼。

在這種情況下，我將構建一個乘法矩陣，該矩陣對我要求和的行具有1s-字典中每個條目的1s集合均不同。

樣本矩陣：

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)