Numba-совместимая реализация np.tile? - PullRequest
Новая
       14

Numba-совместимая реализация np.tile?

0 голосов
/ 08 мая 2020

Я работаю над кодом для удаления дымки изображений на основе этой статьи , и я начал с заброшенной реализации Py2.7 . С тех пор, особенно с Numba, я добился некоторых реальных улучшений производительности (что важно, поскольку мне придется запускать это на изображениях 8K). шаг фильтра коробки (я уже сбрил почти минуту на изображение, но этот последний медленный шаг составляет ~ 30 с / изображение), и я близок к тому, чтобы заставить его работать как nopython в Numba : 

@njit # Row dependencies means can't be parallel
def yCumSum(a):
    """
    Numba based computation of y-direction
    cumulative sum. Can't be parallel!
    """
    out = np.empty_like(a)
    out[0, :] = a[0, :]
    for i in prange(1, a.shape[0]):
        out[i, :] = a[i, :] + out[i - 1, :]
    return out

@njit(parallel= True)
def xCumSum(a):
    """
    Numba-based parallel computation
    of X-direction cumulative sum
    """
    out = np.empty_like(a)
    for i in prange(a.shape[0]):
        out[i, :] = np.cumsum(a[i, :])
    return out

@jit
def _boxFilter(m, r, gpu= hasGPU):
    if gpu:
        m = cp.asnumpy(m)
    out = __boxfilter__(m, r)
    if gpu:
        return cp.asarray(out)
    return out

@jit(fastmath= True)
def __boxfilter__(m, r):
    """
    Fast box filtering implementation, O(1) time.
    Parameters
    ----------
    m:  a 2-D matrix data normalized to [0.0, 1.0]
    r:  radius of the window considered
    Return
    -----------
    The filtered matrix m'.
    """
    #H: height, W: width
    H, W = m.shape
    #the output matrix m'
    mp = np.empty(m.shape)

    #cumulative sum over y axis
    ySum = yCumSum(m) #np.cumsum(m, axis=0)
    #copy the accumulated values of the windows in y
    mp[0:r+1,: ] = ySum[r:(2*r)+1,: ]
    #differences in y axis
    mp[r+1:H-r,: ] = ySum[(2*r)+1:,: ] - ySum[ :H-(2*r)-1,: ]
    mp[(-r):,: ] = np.tile(ySum[-1,: ], (r, 1)) - ySum[H-(2*r)-1:H-r-1,: ]

    #cumulative sum over x axis
    xSum = xCumSum(mp) #np.cumsum(mp, axis=1)
    #copy the accumulated values of the windows in x
    mp[:, 0:r+1] = xSum[:, r:(2*r)+1]
    #difference over x axis
    mp[:, r+1:W-r] = xSum[:, (2*r)+1: ] - xSum[:, :W-(2*r)-1]
    mp[:, -r: ] = np.tile(xSum[:, -1][:, None], (1, r)) - xSum[:, W-(2*r)-1:W-r-1]
    return mp

Там есть чем заняться по краям, но если я смогу получить операцию плитки как вызов python, я не смогу python весь этап boxfilter и получить большой прирост производительности . Я не очень склонен делать что-то действительно конкретное c, так как я бы хотел повторно использовать этот код в другом месте, но я бы не стал особо возражать против его ограничения 2D-областью. По какой-то причине я просто смотрю на это и не совсем уверен, с чего начать.

1 Ответ

1 голос
/ 09 мая 2020

np.tile - это слишком сложный для полной переопределения, но если я не неправильно понимаю, похоже, вам нужно всего лишь взять вектор, а затем повторить его по другой оси r раз .

Numba-совместимый способ сделать это - написать 

y = x.repeat(r).reshape((-1, r))

Затем x будет повторяться r раз по второму измерению, так что y[i, j] == x[i].

Пример: 

In [2]: x = np.arange(5)                                                                                                

In [3]: x.repeat(3).reshape((-1, 3))                                                                                                                                  
Out[3]: 
array([[0, 0, 0],
       [1, 1, 1],
       [2, 2, 2],
       [3, 3, 3],
       [4, 4, 4]])

Если вы хотите, чтобы вместо этого x повторялось по первому измерению, просто возьмите транспонирование y.T.

...