Матричное умножение с итераторной зависимостью - NumPy

Question

Когда-то назад this question (теперь удалено, но 10K + rep пользователи все еще могут просмотреть его). Это выглядело интересно для меня, и я узнал что-то новое там, пытаясь его решить, и я подумал, что стоит поделиться. Я хотел бы опубликовать эти идеи/решения и хотел бы, чтобы люди публиковали другие возможные способы его решения. Я отправляю суть вопроса.

Итак, мы имеем два Numpy ndarrays a и b форм:

a : (m,n,N) 
b : (n,m,N) 

Давайте предположим, что мы имеем дело со случаями, когда m, n & N сопоставимы.

Проблема заключается в том, чтобы решить следующие операции умножения и суммирования с упором на производительность:

def all_loopy(a,b): 
    P,Q,N = a.shape 
    d = np.zeros(N) 
    for i in range(N): 
     for j in range(i): 
      for k in range(P): 
       for n in range(Q): 
        d[i] += a[k,n,i] * b[n,k,j] 
    return d 

Answer 1

я узнал кое-что по пути, пытаясь найти векторизованные и быстрые способы ее решения.

1) Во-первых, существует зависимость итераторов от "for j in range(i)". Из моего предыдущего опыта, особенно с целью решения таких проблем на MATLAB, оказалось, что такую ​​зависимость можно позаботиться с помощью lower triangular matrix, поэтому np.tril должен работать там. Таким образом, полностью Векторизованное решение и не столь памяти эффективное решение (так как он создает промежуточный (N,N) профилированный массив перед окончательным восстановлением в (N,) формой массива) были бы -

def fully_vectorized(a,b): 
    return np.tril(np.einsum('ijk,jil->kl',a,b),-1).sum(1) 

2) Следующее трик/Идея состояла в том, чтобы держать одну петли для итератора i в for i in range(N), но вставьте эту зависимость с индексированием и используя np.einsum для выполнения всех этих умножений и суммирования. Преимуществом будет эффективность памяти. Реализация будет выглядеть так:

def einsum_oneloop(a,b): 
    d = np.zeros(N) 
    for i in range(N): 
     d[i] = np.einsum('ij,jik->',a[:,:,i],b[:,:,np.arange(i)]) 
    return d 

Есть еще два очевидных способа решить эту проблему. Таким образом, если мы начинаем работать с оригинальным all_loopy раствора, можно сохранить внешнюю две петли и использовать np.einsum или np.tensordot для выполнения этих операций и, таким образом, удалить внутреннюю две петли, как так -

def tensordot_twoloop(a,b): 
    d = np.zeros(N) 
    for i in range(N): 
     for j in range(i): 
      d[i] += np.tensordot(a[:,:,i],b[:,:,j], axes=([1,0],[0,1]))   
    return d 

def einsum_twoloop(a,b): 
    d = np.zeros(N) 
    for i in range(N): 
     for j in range(i): 
      d[i] += np.einsum('ij,ji->',a[:,:,i],b[:,:,j]) 
    return d 

время выполнения теста

Давайте сравним все пять подходов, опубликованных до сих пор, чтобы решить проблему, в том числе опубликованную в вопросе.

Дело № 1:

In [26]: # Input arrays with random elements 
    ...: m,n,N = 20,20,20 
    ...: a = np.random.rand(m,n,N) 
    ...: b = np.random.rand(n,m,N) 
    ...: 

In [27]: %timeit all_loopy(a,b) 
    ...: %timeit tensordot_twoloop(a,b) 
    ...: %timeit einsum_twoloop(a,b) 
    ...: %timeit einsum_oneloop(a,b) 
    ...: %timeit fully_vectorized(a,b) 
    ...: 
10 loops, best of 3: 79.6 ms per loop 
100 loops, best of 3: 4.97 ms per loop 
1000 loops, best of 3: 1.66 ms per loop 
1000 loops, best of 3: 585 µs per loop 
1000 loops, best of 3: 684 µs per loop 

Дело № 2:

In [28]: # Input arrays with random elements 
    ...: m,n,N = 50,50,50 
    ...: a = np.random.rand(m,n,N) 
    ...: b = np.random.rand(n,m,N) 
    ...: 

In [29]: %timeit all_loopy(a,b) 
    ...: %timeit tensordot_twoloop(a,b) 
    ...: %timeit einsum_twoloop(a,b) 
    ...: %timeit einsum_oneloop(a,b) 
    ...: %timeit fully_vectorized(a,b) 
    ...: 
1 loops, best of 3: 3.1 s per loop 
10 loops, best of 3: 54.1 ms per loop 
10 loops, best of 3: 26.2 ms per loop 
10 loops, best of 3: 27 ms per loop 
10 loops, best of 3: 23.3 ms per loop 

Дело № 3 (Выход из all_loopy за то, что очень медленно):

In [30]: # Input arrays with random elements 
    ...: m,n,N = 100,100,100 
    ...: a = np.random.rand(m,n,N) 
    ...: b = np.random.rand(n,m,N) 
    ...: 

In [31]: %timeit tensordot_twoloop(a,b) 
    ...: %timeit einsum_twoloop(a,b) 
    ...: %timeit einsum_oneloop(a,b) 
    ...: %timeit fully_vectorized(a,b) 
    ...: 
1 loops, best of 3: 1.08 s per loop 
1 loops, best of 3: 744 ms per loop 
1 loops, best of 3: 568 ms per loop 
1 loops, best of 3: 866 ms per loop 

Идя по номерам , einsum_oneloop выглядит довольно хорошо для меня, тогда как fully_vectorized можно использовать при работе с небольшими приличные массивы!

Answer 2

Я не уверен, что вы хотите, чтобы все было - numpy, но я всегда использовал numba для медленных, но простых в использовании функций на основе петли. Ускорение для сложных циклов задач удивительно. Сначала я просто numba.njit Ted вашего all_loopy варианта, который уже дал мне сравнительную степень результаты:

m,n,N = 20,20,20 
a = np.random.rand(m,n,N) 
b = np.random.rand(n,m,N) 

%timeit numba_all_loopy(a,b) 
1000 loops, best of 3: 476 µs per loop # 3 times faster than everything else 
%timeit tensordot_twoloop(a,b) 
100 loops, best of 3: 16.1 ms per loop 
%timeit einsum_twoloop(a,b) 
100 loops, best of 3: 4.02 ms per loop 
%timeit einsum_oneloop(a,b) 
1000 loops, best of 3: 1.52 ms per loop 
%timeit fully_vectorized(a,b) 
1000 loops, best of 3: 1.67 ms per loop 

Затем я проверил это против вашего 100, 100, 100 случая:

m,n,N = 100,100,100 
a = np.random.rand(m,n,N) 
b = np.random.rand(n,m,N) 

%timeit numba_all_loopy(a,b) 
1 loop, best of 3: 2.35 s per loop 
%timeit tensordot_twoloop(a,b) 
1 loop, best of 3: 3.54 s per loop 
%timeit einsum_twoloop(a,b) 
1 loop, best of 3: 2.58 s per loop 
%timeit einsum_oneloop(a,b) 
1 loop, best of 3: 2.71 s per loop 
%timeit fully_vectorized(a,b) 
1 loop, best of 3: 1.08 s per loop 

Помимо заметив, что мой компьютер намного медленнее, чем у вас - версия numba становится медленнее. Что случилось?

Numpy использует скомпилированные версии, и в зависимости от параметров компилятора numpy, вероятно, оптимизирует цикл, а numba - нет. Итак, следующим логическим шагом является оптимизация цикла. Предполагая, что C-смежные массивы, самые внутренние петли должны быть последней осью массивов. Это самая быстрая изменяющаяся ось, поэтому местонахождение кеша будет лучше.

@nb.njit 
def numba_all_loopy2(a,b): 
    P,Q,N = a.shape 
    d = np.zeros(N) 
    # First axis a, second axis b 
    for k in range(P): 
     # first axis b, second axis a 
     for n in range(Q): 
      # third axis a 
      for i in range(N): 
       # third axis b 
       A = a[k,n,i] # so we have less lookups of the same variable 
       for j in range(i): 
        d[i] += A * b[n,k,j] 
    return d 

так каковы тайминги этой «оптимизированной» функции numba? Можно ли сравнивать с другими или даже избивать их?

m = n = N = 20 
%timeit numba_all_loopy(a,b) 
1000 loops, best of 3: 476 µs per loop 
%timeit numba_all_loopy2(a,b) 
1000 loops, best of 3: 379 µs per loop # New one is a bit faster 

так что немного быстрее для небольших матриц, что о больших из них:

m = n = N = 100 
%timeit numba_all_loopy(a,b) 
1 loop, best of 3: 2.34 s per loop 
%timeit numba_all_loopy2(a,b) 
1 loop, best of 3: 213 ms per loop # More then ten times faster now! 

Таким образом, мы имеем удивительное ускорение для больших массивов. Эта функция теперь в 4-5 раз быстрее, чем ваши векторизованные подходы, и результат один и тот же.

Но удивительно кажется, что заказ кажется каким-то образом зависящим от компьютера, потому что fully_vectorized является самым быстрым, где einsum -выключаются быстрее на машине Дивакара. Поэтому он может быть открытым, если эти результаты на самом деле намного быстрее.

Просто для удовольствия я попробовал с n=m=N=200:

%timeit numba_all_loopy2(a,b) 
1 loop, best of 3: 3.38 s per loop # still 5 times faster 
%timeit einsum_oneloop(a,b) 
1 loop, best of 3: 51.4 s per loop 
%timeit fully_vectorized(a,b) 
1 loop, best of 3: 16.7 s per loop