有时,如果您真的想加快计算速度,而原生numpy无法做到,则需要编写非惯用的numpy代码。
numba
将您的python代码编译为低级C。由于许多numpy本身通常与C一样快,所以如果您的问题不适合使用numpy进行本机矢量化,那么这最终将非常有用。这是一个示例(我假设索引是连续的并且已排序,这也反映在示例数据中):
import numpy as np
import numba
# use the inflated example of roganjosh https://stackoverflow.com/a/58788534
data = [1.00, 1.05, 1.30, 1.20, 1.06, 1.54, 1.33, 1.87, 1.67]
index = [0, 0, 1, 1, 1, 1, 2, 3, 3]
data = np.array(data * 500) # using arrays is important for numba!
index = np.sort(np.random.randint(0, 30, 4500))
# jit-decorate; original is available as .py_func attribute
@numba.njit('f8[:](f8[:], i8[:])') # explicit signature implies ahead-of-time compile
def diffmedian_jit(data, index):
res = np.empty_like(data)
i_start = 0
for i in range(1, index.size):
if index[i] == index[i_start]:
continue
# here: i is the first _next_ index
inds = slice(i_start, i) # i_start:i slice
res[inds] = data[inds] - np.median(data[inds])
i_start = i
# also fix last label
res[i_start:] = data[i_start:] - np.median(data[i_start:])
return res
以下是一些使用IPython的%timeit
魔术的时机:
>>> %timeit diffmedian_jit.py_func(data, index) # non-jitted function
... %timeit diffmedian_jit(data, index) # jitted function
...
4.27 ms ± 109 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
65.2 µs ± 1.01 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
使用问题中的更新示例数据,这些数字(即python函数的运行时与JIT加速功能的运行时)是
>>> %timeit diffmedian_jit.py_func(data, groups)
... %timeit diffmedian_jit(data, groups)
2.45 s ± 34.4 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
93.6 ms ± 518 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
使用加速代码,这在较小的情况下相当于65倍的加速,在较大的情况下相当于26倍的加速(当然,与慢速循环代码相比)。另一个好处是(与使用本地numpy进行典型矢量化处理不同),我们不需要额外的内存来达到这种速度,而这全都在于优化和编译的低级代码最终得以运行。
上面的函数假定int64
默认情况下为numpy int数组,而在Windows上实际上并非如此。因此,另一种方法是从的调用中删除签名,以numba.njit
触发适当的即时编译。但这意味着函数将在第一次执行时进行编译,这可能会与计时结果混为一谈(我们可以使用代表性的数据类型手动执行一次该函数,或者只接受第一次计时执行会慢得多,这应该被忽略)。这正是我试图通过指定签名来阻止提前触发编译的尝试。
无论如何,在适当的JIT情况下,我们需要的装饰器只是
@numba.njit
def diffmedian_jit(...):
请注意,我为jit编译的函数显示的上述时序仅在编译函数后才适用。这要么在定义时发生(使用急切编译,将显式签名传递给numba.njit
),要么在第一个函数调用期间发生(使用惰性编译,当没有签名传递给numba.njit
)。如果该函数仅执行一次,则该方法的速度也应考虑编译时间。通常,只有在编译+执行的总时间少于未编译的运行时的情况下,才值得编译函数(在上述情况下,这是正确的,因为本机python函数非常慢)。这通常发生在您多次调用编译函数时。
正如max9111在评论中指出的,的一个重要功能numba
是cache
关键字 to jit
。传递cache=True
给numba.jit
会将已编译的函数存储到磁盘,以便在给定python模块的下一次执行期间,该函数将从那里加载而不是重新编译,这从长远来看又可以节省您的运行时间。
scipy.ndimage.median
在链接答案中添加了建议时间?在我看来,每个标签不需要相等数量的元素。还是我错过了什么?