Много раз повторять np.einsum ... Есть ли более быстрый способ?

Question

У меня есть функция правдоподобия, которую я пытаюсь сэмплировать с помощью MCM C. Я использовал цикл no for в самом журнале вероятности, но я вызываю np.einsum() один раз.

Вот пример того, как выглядит мой текущий код:

A = np.random.rand(4,50,60,200) # Random NDarray
B = np.random.rand(200,1000,4)  # Random NDarray
out = np.einsum('ijkl,lui->jkui', A, B, optimize="optimal")

Результат out имеет размеры (50,60,1000,4). Этот расчет слишком медленный, чтобы обеспечить эффективную выборку MCM C (~ 4 секунды на моей машине), есть ли способ ускорить его? Одна полезная информация заключается в том, что при каждом вызове функции логарифма правдоподобия, в то время как фактические значения в массивах A и B меняются, размеры каждого массива остаются фиксированными. Я полагаю, что это могло бы быть полезным для ускорения работы, поскольку одни и те же элементы всегда умножаются вместе.

Answer 1

Ну, одна из осей остается выровненной в A (первая) и B (последняя), а также остается на выходе (последняя) и представляет собой очень маленькое число цикла 4. Итак, мы могли бы просто l oop вместо этого с np.tensordot для тензорного уменьшения суммы. Преимущество 4x меньшей перегрузки памяти при работе с такими большими наборами данных могло бы преодолеть 4-кратное зацикливание, потому что вычисление на итерацию также 4x меньше.

Таким образом, решение с tensordot будет -

def func1(A, B):
    out = np.empty(A.shape[1:3] + B.shape[1:])
    for i in range(len(A)):
        out[...,i] = np.tensordot(A[i], B[...,i],axes=(-1,0))
    return out

Тайминги -

In [70]: A = np.random.rand(4,50,60,200) # Random NDarray
    ...: B = np.random.rand(200,1000,4)  # Random NDarray
    ...: out = np.einsum('ijkl,lui->jkui', A, B, optimize="optimal")

# Einsum solution without optimize    
In [71]: %timeit np.einsum('ijkl,lui->jkui', A, B)
2.89 s ± 109 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

# Einsum solution with optimize    
In [72]: %timeit np.einsum('ijkl,lui->jkui', A, B, optimize="optimal")
2.79 s ± 9.31 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)    

# @Paul Panzer's soln
In [74]: %timeit np.stack([np.tensordot(a,b,1) for a,b in zip(A,B.transpose(2,0,1))],-1)
183 ms ± 6.08 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

In [73]: %timeit func1(A,B)
158 ms ± 3.35 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

Чтобы еще раз повторить важность перегрузки памяти и требований к вычислениям, допустим, мы хотим суммировать-уменьшить последнюю ось длины 4 Кроме того, тогда мы увидим заметную разницу во времени для версии optimal -

In [78]: %timeit np.einsum('ijkl,lui->jkui', A, B, optimize="optimal")
2.76 s ± 9.36 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [79]: %timeit np.einsum('ijkl,lui->jku', A, B, optimize="optimal")
93.8 ms ± 3.3 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

Так что в этом случае было бы лучше go с einsum.

Спецификация c для данной проблемы

Учитывая, что размеры A и B остаются неизменными, инициализация массива с помощью out = np.empty(A.shape[1:3] + B.shape[1:]) может быть выполнена как одноразовое дело и l oop при каждом вызове функции логарифма правдоподобия с предложенным циклом для использования tensordot и обновления вывода out.

Answer 2

Даже при использовании в маленьком l oop tensordot более чем в 10 раз быстрее:

timeit(lambda:np.einsum('ijkl,lui->jkui', A, B, optimize="optimal"),number=5)/5
# 3.052245747600682
timeit(lambda:np.stack([np.tensordot(a,b,1) for a,b in zip(A,B.transpose(2,0,1))],-1),number=10)/10
# 0.23842503569903784

out_td = np.stack([np.tensordot(a,b,1) for a,b in zip(A,B.transpose(2,0,1))],-1)
out_es = np.einsum('ijkl,lui->jkui', A, B, optimize="optimal")
np.allclose(out_td,out_es)
# True