Numba在頻率計數方面比純Python慢

Question

給定一個數據矩陣，其中離散條目表示為2D numpy數組，我試圖計算一些特征（列）的觀察頻率，僅查看某些實例（矩陣的行）。

在完成一些奇特的切片之后，我可以很容易地使用bincount應用於每個切片的numpy。 使用外部數據結構作為計數累加器，在純Python中執行此操作是C風格的雙循環。

import numpy

import numba

try:
    from time import perf_counter
except:
    from time import time
    perf_counter = time


def estimate_counts_numpy(data,
                          instance_ids,
                          feature_ids):
    """
    WRITEME
    """
    #
    # slicing the data array (probably memory consuming)
    curr_data_slice = data[instance_ids, :][:, feature_ids]

    estimated_counts = []
    for feature_slice in curr_data_slice.T:
        counts = numpy.bincount(feature_slice)
        #
        # checking just for the all 0 case:
        # this is not stable for not binary datasets TODO: fix it
        if counts.shape[0] < 2:
            counts = numpy.append(counts, [0], 0)
        estimated_counts.append(counts)

    return estimated_counts


@numba.jit(numba.types.int32[:, :](numba.types.int8[:, :],
                                   numba.types.int32[:],
                                   numba.types.int32[:],
                                   numba.types.int32[:],
                                   numba.types.int32[:, :]))
def estimate_counts_numba(data,
                          instance_ids,
                          feature_ids,
                          feature_vals,
                          estimated_counts):
    """
    WRITEME
    """

    #
    # actual counting
    for i, feature_id in enumerate(feature_ids):
        for instance_id in instance_ids:
            estimated_counts[i][data[instance_id, feature_id]] += 1

    return estimated_counts


if __name__ == '__main__':
    #
    # creating a large synthetic matrix, testing for performance
    rand_gen = numpy.random.RandomState(1337)
    n_instances = 2000
    n_features = 2000
    large_matrix = rand_gen.binomial(1, 0.5, (n_instances, n_features))
    #
    # random indexes too
    n_sample = 1000
    rand_instance_ids = rand_gen.choice(n_instances, n_sample, replace=False)
    rand_feature_ids = rand_gen.choice(n_features, n_sample, replace=False)
    binary_feature_vals = [2 for i in range(n_features)]
    #
    # testing
    numpy_start_t = perf_counter()

    e_counts_numpy = estimate_counts_numpy(large_matrix,
                                           rand_instance_ids,
                                           rand_feature_ids)
    numpy_end_t = perf_counter()
    print('numpy done in {0} secs'.format(numpy_end_t - numpy_start_t))

    binary_feature_vals = numpy.array(binary_feature_vals)
    #
    #
    curr_feature_vals = binary_feature_vals[rand_feature_ids]
    #
    # creating a data structure to hold the slices
    # (with numba I cannot use list comprehension?)
    # e_counts_numba = [[0 for val in range(feature_val)]
    #                   for feature_val in
    #                   curr_feature_vals]
    e_counts_numba = numpy.zeros((n_sample, 2), dtype='int32')
    numba_start_t = perf_counter()

    estimate_counts_numba(large_matrix,
                          rand_instance_ids,
                          rand_feature_ids,
                          binary_feature_vals,
                          e_counts_numba)
    numba_end_t = perf_counter()
    print('numba done in {0} secs'.format(numba_end_t - numba_start_t))

這些是我在運行上面代碼時得到的時間：

numpy done in 0.2863295429997379 secs
numba done in 11.55551904299864 secs

我的觀點是，當我嘗試使用numba應用jit時，我的實現速度更慢，所以我非常懷疑我搞砸了。

Answer 1

你的函數很慢的原因是因為Numba已經回到對象模式來編譯循環。

有兩個問題：

1. Numba尚不支持多維數組的鏈式索引，因此您需要重寫：

estimated_counts[i][data[instance_id, feature_id]]

進入這個：

estimated_counts[i, data[instance_id, feature_id]]

2. 您的顯式類型簽名不正確。 所有輸入數組實際上都是int64，而不是int8 / int32。 您可以依靠Numba的自動JIT來檢測參數類型並編譯正確的版本，而不是修復您的簽名。 您所要做的就是將裝飾器更改為@numba.jit 。 如果您不想包含編譯時間，請確保在基准測試之前調用該函數一次。

通過這些更改，我將Numba的基准測試比NumPy快15％。