Блог

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

Сейчас я ищу способ использовать многопроцессорную обработку для ускорения rebin функции. В настоящее время функция использует один процессор. У меня есть некоторые знания о модуле многопроцессорной обработки, но я был бы признателен за совет по преобразованию приведенного ниже кода для использования многопроцессорной обработки.

ПРИМЕЧАНИЕ: я ценю рекомендации (возможно, более питоновские), которые могут ускорить приведенную ниже функцию по-другому. Тем не менее, я также хотел бы узнать, может ли кто-нибудь помочь обновить приведенный ниже код, чтобы он подходил для многопроцессорной обработки.

MWE:

 import numpy as np
from tqdm.auto import tqdm

def rebin(arr, new_shape):
    shape = (arr.shape[0], new_shape[0], arr.shape[1] // new_shape[0],
             new_shape[1], arr.shape[2] // new_shape[1])
    return arr.reshape(shape).mean(-1).mean(2)


def rebinner(arr, window):
    # w is the window size used to perform rebin on a fixed size window as the entire
    # data cannot be fit in memory for large data size.
    y = []
    for j in tqdm(range(0, arr.shape[0], window)):
        y.append(rebin(arr[j:j   window, :, None]*arr[j:j   window, None, :], [32, 32]))
    return np.concatenate(y, axis=0)


arr = np.random.random((1000,2560,))
  

Приведенный ниже код выполняется в ячейке jupyter notebook для проверки времени выполнения. Также может использоваться в скрипте с использованием timeit.

 %%time
print(rebinner(arr, window=10).shape)
  

Ожидаемый результат:

 (1000, 32, 32)
CPU times: user 8.52 s, sys: 4.58 s, total: 13.1 s
Wall time: 13.1 s
  

ОБНОВЛЕНИЕ 1:

Вывод с использованием библиотеки numba, предложенной @john-zwinck

Основываясь на комментариях от @johnzwinck, я немного обновил код и включил декораторы numa. Однако новый скрипт выдает ошибку утверждения, и я не совсем уверен, что является причиной этого. Ниже приведен обновленный код и соответствующее сообщение об ошибке.

 import numba
import numpy as np


@numba.njit(nopython=True)
def rebinner2(arr, new_shape):
    shape = (arr.shape[0], new_shape[0], arr.shape[1] // new_shape[0],
             new_shape[1], arr.shape[2] // new_shape[1])
    return arr.reshape(shape).mean(-1).mean(2)


@numba.njit(nopython=True)
def rebinner1(arr, window):
    return [rebinner2(np.random.random((window, 1, 2560))*np.random.random((window, 2560, 1)),
                      [32, 32]) for j in range(0, arr.shape[0], window)]
    # return [rebinner2(arr[j:j   window, :, None]*arr[j:j   window, None, :],
    # [32, 32]) for j in range(0, arr.shape[0], window)]


def rebinner(arr, window):
    return np.concatenate(rebinner1(arr,window), axis=0)


if __name__ == "__main__":
    arr = np.random.random((1000, 2560))
    print(rebinner(arr, window=10).shape)
    # print(rebinner1(arr, window=10).shape)
  

Output:

 venv/lib/python3.6/site-packages/numba/core/decorators.py:252: RuntimeWarning: nopython is set for njit and is ignored
  warnings.warn('nopython is set for njit and is ignored', RuntimeWarning)
Traceback (most recent call last):
  File "numba_tester.py", line 28, in <module>
    print(rebinner(arr, window=10).shape)
  File "numba_tester.py", line 23, in rebinner
    return np.concatenate(rebinner1(arr,window), axis=0)
  File "venv/lib/python3.6/site-packages/numba/core/dispatcher.py", line 415, in _compile_for_args
    error_rewrite(e, 'typing')
  File "venv/lib/python3.6/site-packages/numba/core/dispatcher.py", line 358, in error_rewrite
    reraise(type(e), e, None)
  File "venv/lib/python3.6/site-packages/numba/core/utils.py", line 80, in reraise
    raise value.with_traceback(tb)
numba.core.errors.TypingError: Failed in nopython mode pipeline (step: nopython frontend)
Failed in nopython mode pipeline (step: nopython frontend)
- Resolution failure for literal arguments:
AssertionError()
- Resolution failure for non-literal arguments:
AssertionError()

During: resolving callee type: BoundFunction(array.mean for array(float64, 5d, C))
During: typing of call at numba_tester.py (9)


File "numba_tester.py", line 9:
def rebinner2(arr, new_shape):
    <source elided>
             new_shape[1], arr.shape[2] // new_shape[1])
    return arr.reshape(shape).mean(-1).mean(2)
    ^

During: resolving callee type: type(CPUDispatcher(<function rebinner2 at 0x7f6c5743e730>))
During: typing of call at numba_tester.py (17)

During: resolving callee type: type(CPUDispatcher(<function rebinner2 at 0x7f6c5743e730>))
During: typing of call at numba_tester.py (17)


File "numba_tester.py", line 17:
def rebinner1(arr, window):
    <source elided>
    return [rebinner2(np.random.random((window, 1, 2560))*np.random.random((window, 2560, 1)),
                      [32, 32]) for j in range(0, arr.shape[0], window)]
  

Примечания: Приведенный выше код отлично работает при реализации без декораторов numba. Я упростил rebinner2 , используя случайные массивы вместо фактической операции среза-транспонирования-умножения (чтобы проверить, не вызвало ли это проблем для numba, но это не так).

ОБНОВЛЕНИЕ 2:

Модифицированный код numba работает, но прироста производительности нет.

По предложению @JohnZwinck я изменил njit на jit . Это отключает nopython режим. Однако при проверке времени, необходимого для выполнения, кажется, что метод numba теперь работает хуже (вероятно, потому, что отключение nopython теряет оптимизацию).

С помощью Numba:

 (5000, 32, 32)
real:   1m11.538s
user:   0m49.137s
sys :   0m22.872s
  

Без Numba:

 (5000, 32, 32)
real:   1m2.439s
user:   0m41.721s
sys :   0m21.152s
  

Просто удалите ненужное tqdm (индикатор выполнения) и добавьте Numba, чтобы ускорить его вот так (непроверенный, но должен быть близок к тому, что вам нужно):

 import numba

@numba.njit
def rebinner2(arr, new_shape):
    shape = (arr.shape[0], new_shape[0], arr.shape[1] // new_shape[0],
             new_shape[1], arr.shape[2] // new_shape[1])
    return arr.reshape(shape).mean(-1).mean(2)

@numba.njit
def rebinner1(arr, window):
    return [rebinner2(arr[j:j   w, :, None]*arr[j:j   w, None, :], [32, 32]) for j in range(0, arr.shape[0], window)]

def rebinner(arr, window):
    return np.concatenate(rebinner1, axis=0)
  

Я завернул np.concatenate() в функцию, отличную от Numba, потому что я не уверен, поддерживает ли Numba axis аргумент.

Использование нескольких ядер для этого, вероятно, бессмысленно, если ваши данные действительно гигантские, и в этом случае вам следует просто загружать подмножество данных в каждый процесс. Но Numba сделает цикл намного быстрее, чем несколько ядер.