Question

我有一个代码，首先我需要对值进行排序，然后我需要对前10个元素求和。我很想使用Numba软件包来加快运行时间，但它不起作用，Numba的代码比Numpy慢。

我的第一次测试，只是为了总和：

import numpy as np
import numba
np.random.seed(0)

def SumNumpy(x):
    return np.sum(x[:10])

@numba.jit()
def SumNumpyNumba(x):
    return np.sum(x[:10])

我的测试：

x = np.random.rand(1000000000)
%timeit SumNumpy(x)
%timeit SumNumpyNumba(x)

结果：

100000个循环，最佳3：每循环6.8μs

1000000次循环，最佳3：715 ns /循环

这没关系，Numba做得很好。但是当我一起尝试np.sort和np.sum时：

def sumSortNumpy(x):
    y = np.sort(x)
    return np.sum(y[:10])

@numba.jit()
def sumSortNumpyNumba(x):
    y = np.sort(x)
    return np.sum(y[:10])

并测试：

x = np.random.rand(100000)
%timeit sumSortNumpy(x)
%timeit sumSortNumpyNumba(x)

结果：

100个循环，最佳3：每循环14.6毫秒

10个循环，最佳3：每循环20.6 ms

Numba / Numpy变得比Numpy慢。所以我的问题是否有什么可以改善功能“sumSortNumpyNumba”？

我感谢你的帮助。

感谢。

Answer 1

我们在排序后进行求和，因此顺序在第一个N=10元素中无关紧要。因此，我们可以使用np.argpartition来避免排序步骤，并简单地为我们提供可以在以后求和的第一个N最小数字的组，就像这样 -

def sumSortNumPyArgpartition(x, N=10):
    return x[np.argpartition(x, N)[:N]].sum()

各种数据集上的计时 -

In [39]: np.random.seed(0)
    ...: x = np.random.rand(1000000)

In [40]: %timeit sumSortNumpy(x)
    ...: %timeit sumSortNumPyArgpartition(x)
10 loops, best of 3: 78.6 ms per loop
100 loops, best of 3: 12.3 ms per loop

In [41]: np.random.seed(0)
    ...: x = np.random.rand(10000000)

In [42]: %timeit sumSortNumpy(x)
    ...: %timeit sumSortNumPyArgpartition(x)
1 loop, best of 3: 920 ms per loop
10 loops, best of 3: 153 ms per loop

In [43]: np.random.seed(0)
    ...: x = np.random.rand(100000000)

In [44]: %timeit sumSortNumpy(x)
    ...: %timeit sumSortNumPyArgpartition(x)
1 loop, best of 3: 10.6 s per loop
1 loop, best of 3: 978 ms per loop

有效地计算数组

1 个答案: