Numba на вложенных массивах Numpy - PullRequest
Новая
       18

Numba на вложенных массивах Numpy

2 голосов
/ 11 июля 2019

Настройка

У меня есть следующие две реализации матричного вычисления:

  1. Первая реализация использует matrix of shape (n, m) и вычисление повторяется в цикле for для repetition -times: 
import numpy as np
from numba import jit

@jit
def foo():
    for i in range(1, n):
        for j in range(1, m):

            _deleteA = (
                        matrix[i, j] +
                        #some constants added here
            )
            _deleteB = (
                        matrix[i, j-1] +
                        #some constants added here
            )
            matrix[i, j] = min(_deleteA, _deleteB)

    return matrix

repetition = 3
for x in range(repetition):
    foo()


2. Вторая реализация исключает дополнительный цикл for и, следовательно, включает в себя repetition = 3 в матрицу, которая тогда имеет значение shape (repetition, n, m): 

@jit
def foo():
    for i in range(1, n):
        for j in range(1, m):

            _deleteA = (
                        matrix[:, i, j] +
                        #some constants added here
            )
            _deleteB = (
                        matrix[:, i, j-1] +
                        #some constants added here
            )
            matrix[:, i, j] = np.amin(np.stack((_deleteA, _deleteB), axis=1), axis=1)

    return matrix


Вопросы

Относительно обеих реализаций, яобнаружил две вещи относительно их производительности с %timeit в iPython.

  1. Первая реализация сильно выигрывает от @jit, а вторая - вовсе (28 мс против 25 с в моем тестовом примере). Кто-нибудь может представить, почему @jit больше не работает с массивом фигур (repetition, n, m)?


Редактировать

Я переместил предыдущий второй вопрос в дополнительный пост , поскольку задание нескольких вопросов считается плохим стилем SO.

Вопросбыло:

При игнорировании @jit первая реализация все еще намного быстрее (тот же тест: 17 с против 26 с). Почему numpy медленнее при работе с тремя, а не с двумя измерениями?

1 Ответ

3 голосов
/ 11 июля 2019

Я не уверен, какие у вас настройки, но я немного переписал ваш пример: 

import numpy as np
from numba import jit

#@jit(nopython=True)
def foo(matrix):
    n, m = matrix.shape
    for i in range(1, n):
        for j in range(1, m):

            _deleteA = (
                        matrix[i, j] #+
                        #some constants added here
            )
            _deleteB = (
                        matrix[i, j-1] #+
                        #some constants added here
            )
            matrix[i, j] = min(_deleteA, _deleteB)

    return matrix

foo_jit = jit(nopython=True)(foo)

, а затем время: 

m = np.random.normal(size=(100,50))

%timeit foo(m)  # in a jupyter notebook
# 2.84 ms ± 54.2 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

%timeit foo_jit(m)  # in a jupyter notebook
# 3.18 µs ± 38.9 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

Так что здесь numba - этонамного быстрее, чем ожидалось.Следует учитывать, что глобальные массивы с номерами не ведут себя в numba так, как вы могли бы ожидать:

https://numba.pydata.org/numba-doc/dev/user/faq.html#numba-doesn-t-seem-to-care-when-i-modify-a-global-variable

Обычно лучше передавать данные, как я делал в примере.

Ваша проблема во втором случае заключается в том, что numba не поддерживает amin в настоящее время.См .:

https://numba.pydata.org/numba-doc/dev/reference/numpysupported.html

Вы можете увидеть это, если передадите nopython=True на jit.Таким образом, в текущих версиях numba (0.44 или более ранней версии в настоящее время) он будет возвращаться к objectmode, что часто не быстрее, чем без использования numba, а иногда медленнее, поскольку возникают некоторые накладные расходы на вызовы.

...