非常大的 numpy 数组中更快的中位数

Question

我有一个非常大的 numpy 数组，维度为 (4000, 6000, 15)。

我现在想要每个堆栈的中位数，即沿着第三个维度。当前代码有效，但出奇地慢，单个堆栈 [0,0,:]（15 个值）的中位数至少需要半秒左右才能完成。

height = 4000
width = 6000
N = 15

poolmedian = np.zeros((height,width,3))
RGBmedian = np.zeros((height,width,N), dtype=float)    

for n in range(0,height):
    for m in range(0,width):
                poolmedian[n,m,0] = np.median(RGBmedian[n,m,:])

Answer 1

您需要尽可能地向量化中值计算。每次调用 numpy 函数时，都会在 C 层和 Python 层之间来回移动。 C层尽量多做：

import numpy as np
height = 40
width = 60
N = 15

np.random.seed(1)
poolmedian = np.zeros((height,width,3))
RGBmedian = np.random.random((height,width,N))

def original():
    for n in range(0,height):
        for m in range(0,width):
            poolmedian[n,m,0] = np.median(RGBmedian[n,m,:])
    return poolmedian

def vectorized():
    # Note: np.median is only called ONCE, not n*m times.
    poolmedian[:, :, 0] = np.median(RGBmedian, axis=-1)
    return poolmedian


orig = original()
vec = vectorized()

np.testing.assert_array_equal(orig, vec)

您可以看到自断言通过以来值是相同的（尽管不清楚为什么在 poolmedian 中需要 3 个 dims）。我将上面的代码放在一个名为 test.py 的文件中，并使用 IPython 因为它很方便 %timeit。我还稍微降低了大小，以便它运行得更快，但是您应该在大数据上获得类似的节省。矢量化版本大约快 100 倍：

In [1]: from test import original, vectorized

In [2]: %timeit original()
69.1 ms ± 394 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

In [3]: %timeit vectorized()
618 µs ± 4.1 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

一般来说，您希望使用 numpys 广播规则并尽可能少地调用一个函数。如果您正在寻找高性能 numpy 代码，那么在循环中调用函数几乎总是一个禁忌。

附录：

我已经在test.py中添加了下面的函数，因为还有另一个答案，我想说清楚调用完全矢量化版本（即没有循环）更快，并且修改为使用 4000 x 6000 调光的代码：

import numpy as np
height = 4000
width = 6000
N = 15

...

def fordy():
    for n in range(0,height):
        for m in range(0,width):
            array = RGBmedian[n,m,:]
            array.sort()
            poolmedian[n, m, 0] = (array[6] + array[7])/2
    return poolmedian

如果我们将所有这些加载到 IPython，我们将得到：

In [1]: from test import original, fordy, vectorized

In [2]: %timeit original()
6.87 s ± 72.6 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [3]: %timeit fordy()
262 ms ± 737 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [4]: %timeit vectorized()
18.4 ms ± 149 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

HTH.

非常大的 numpy 数组中更快的中位数

Faster median in very large numpy arrays

python

numpy

median

large-data