Быстрее медиана в очень больших массивах NumPy - PullRequest
Новая
       12

Быстрее медиана в очень больших массивах NumPy

0 голосов
/ 25 октября 2018

У меня есть очень большой массив с размером (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,:])

1 Ответ

0 голосов
/ 25 октября 2018

Вы хотите максимально векторизовать вычисление медианы.Каждый раз, когда вы вызываете функцию 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)

Вы можете видеть, что значения одинаковы после прохождения утверждения (хотя не ясно, зачем вам нужно 3 димса в poolmedian).Я поместил приведенный выше код в файл с именем 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)

В общем, вы хотите использовать правила вещания numpy s и вызывать функцию как можно меньше раз.Вызов функций в цикле - это почти всегда нет-нет, если вам нужен код совершенного numpy.

Приложение:

Я добавил следующую функцию в test.py, так как есть другой ответ, я хочу прояснить, что быстрее вызывать полностью векторизованную версию (т.е. без циклов), а также изменить код для использования dims 4000 на 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.

...