Есть ли шанс сделать это быстрее? (numpy .einsum)

Question

Я пытаюсь умножить три массива (A x B x A) с размерами (19000, 3) x (19000, 3, 3) x (19000, 3), так что в конце я получаю 1d-массив размером (19000), поэтому я хочу умножать только по последним одному / двум измерениям.

У меня он работает с np.einsum (), но мне интересно, есть ли там есть способ сделать это быстрее, так как это узкое место всего моего кода.

np.einsum('...i,...ij,...j', A, B, A)

Я уже пробовал это с двумя отдельными вызовами np.einsum (), но что дало мне такую ​​же производительность:

np.einsum('...i, ...i', np.einsum('...i,...ij', A, B), A)

Я уже пробовал оператор @ и добавлял несколько дополнительных осей, но это тоже не ускорило его:

(A[:, None]@B@A[...,None]).squeeze()

Я пытался заставить его работать с np.inner (), np.dot (), np.tensordot () и np.vdot (), но они никогда не давали у меня те же результаты, поэтому я не мог их сравнить.

Есть другие идеи? Есть ли способ повысить производительность?

Я уже бегло просмотрел Numba, но поскольку Numba не поддерживает np.einsum () и многие другие NumPy функции, я бы придется переписать много кода.

Answer 1

Вы можете использовать Numba

Вначале всегда полезно посмотреть, что делает np.einsum. С optimize==optimal обычно действительно хорошо найти способ сокращения, который имеет меньше FLOP. В этом случае возможна лишь незначительная оптимизация, а промежуточный массив относительно велик (я буду придерживаться наивной версии). Также следует отметить, что схватывания с очень маленькими (фиксированными?) Размерами - это особый случай. Это также причина того, почему здесь довольно легко превзойти np.einsum (развернуть et c ..., что делает компилятор, если знает, что al oop состоит только из 3 элементов)

import numpy as np

A=np.random.rand(19000, 3)
B=np.random.rand(19000, 3, 3)

print(np.einsum_path('...i,...ij,...j', A, B, A,optimize="optimal")[1])

"""
  Complete contraction:  si,sij,sj->s
         Naive scaling:  3
     Optimized scaling:  3
      Naive FLOP count:  5.130e+05
  Optimized FLOP count:  4.560e+05
   Theoretical speedup:  1.125
  Largest intermediate:  5.700e+04 elements
--------------------------------------------------------------------------
scaling                  current                                remaining
--------------------------------------------------------------------------
   3                  sij,si->js                                 sj,js->s
   2                    js,sj->s                                     s->s

"""

Реализация Numba

import numba as nb

#si,sij,sj->s
@nb.njit(fastmath=True,parallel=True,cache=True)
def nb_einsum(A,B):
    #check the input's at the beginning
    #I assume that the asserted shapes are always constant
    #This makes it easier for the compiler to optimize 
    assert A.shape[1]==3
    assert B.shape[1]==3
    assert B.shape[2]==3

    #allocate output
    res=np.empty(A.shape[0],dtype=A.dtype)

    for s in nb.prange(A.shape[0]):
        #Using a syntax like that is also important for performance
        acc=0
        for i in range(3):
            for j in range(3):
                acc+=A[s,i]*B[s,i,j]*A[s,j]
        res[s]=acc
    return res

Время

#warmup the first call is always slower 
#(due to compilation or loading the cached function)
res=nb_einsum(A,B)
%timeit nb_einsum(A,B)
#43.2 µs ± 1.22 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
%timeit np.einsum('...i,...ij,...j', A, B, A,optimize=True)
#450 µs ± 8.28 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
%timeit np.einsum('...i,...ij,...j', A, B, A)
#977 µs ± 4.14 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
np.allclose(np.einsum('...i,...ij,...j', A, B, A,optimize=True),nb_einsum(A,B))
#True