Как ускорить numpy тензор * тензорную операцию - PullRequest
Новая
       4

Как ускорить numpy тензор * тензорную операцию

2 голосов
/ 24 января 2020

У меня узкое место в моем коде, это numpy умножение трехмерного массива на * operator с numpy трехмерным массивом.
Я хотел ускорить эту часть программы с помощью декоратора numba @njit или @jit , но это снизило производительность в 2 раза.
Медленная часть кода: 

@numba.jit
def mat_mul_and_sum(img1, img2, alpha):
    return img1*(1-alpha) + img2*alpha

img1, img2 и alpha - это трехмерные массивы np.array с одинаковыми формами.
Возможно ли ускорить эту строку кода?

Ответы [ 2 ]

4 голосов
/ 24 января 2020

Одна опция на самом деле использует numba так, как она должна использоваться (а не просто применяет декоратор). Однако для вашей конкретной функции вы можете использовать многоядерный рендеринг, используя пакет numexpr. 

import numexpr as ne

def mat_mul_and_sum_numexpr(a, b, alpha):
    return ne.evaluate('a*(1-alpha) + b*alpha')

Используя время из другого ответа: 

In [11]: %timeit mat_mul_and_sum(img1, img2, alpha)
21.6 ms ± 955 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

In [12]: %timeit mat_mul_and_sum2(img1, img2, alpha)
6.35 ms ± 126 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [13]: %timeit mat_mul_and_sum_numexpr(img1, img2, alpha)
4.22 ms ± 54.5 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [14]: np.allclose(mat_mul_and_sum(img1, img2, alpha), mat_mul_and_sum_numexpr(img1, img2, alpha))
Out[14]: True

Возможно, вам удастся добиться некоторой дополнительной производительности с распараллеливанием numba, но часто использование numexpr обеспечивает хороший прирост производительности без необходимости переписывания кода.

3 голосов
/ 24 января 2020

Если развернуть циклы следующим образом, для размера массива (100, 100, 100) numba будет в два раза быстрее, чем чистая версия numpy, вероятно, из-за того, что не нужно выделять промежуточные массивы : 

import numpy as np
import numba as nb

def mat_mul_and_sum(img1, img2, alpha):
    return img1*(1-alpha) + img2*alpha


@nb.jit
def mat_mul_and_sum2(img1, img2, alpha):
    NI, NJ, NK = img1.shape
    out = np.empty((NI, NJ, NK))

    for i in range(NI):
        for j in range(NJ):
            for k in range(NK):
                out[i,j,k] = img1[i,j,k] * (1.0 - alpha[i,j,k]) + img2[i,j,k] * alpha[i,j,k]

    return out

и затем тестирование: 

N = 100
img1 = np.random.normal(size=(N, N, N))
img2 = np.random.normal(size=(N, N, N))
alpha = np.random.normal(size=(N, N, N))

A = mat_mul_and_sum(img1, img2, alpha)
B = mat_mul_and_sum2(img1, img2, alpha)

np.allclose(A,B) #True

%timeit mat_mul_and_sum(img1, img2, alpha)
# 4.6 ms ± 44.2 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
%timeit mat_mul_and_sum2(img1, img2, alpha)
# 2.4 ms ± 10.5 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

Обновление: Вы также можете попробовать изменить декоратор на nb.jit(parallel=True), а затем заменить внешний l oop с for i in nb.prange(NI):, который на моей машине сбрасывает результаты с timeit до 1,37 мс. Это и другие значения времени, безусловно, будут отличаться от машины к машине, а также от размера входов.

...