Оптимизация вызовов np.einsum в Python
У меня есть два пустых массива: X размера (N,N,N,N) и Y размера (N,N). Моя цель — как можно быстрее оценить следующий вызов einsum:
Z = np.einsum('iiii,ij,ik,il,im->jklm', X, Y, Y, Y, Y)
Для этого я протестировал три реализации вышеперечисленного, которые дают одинаковые ответы, но по-разному оптимизированы:
def einsum_one(X, Y):
return np.einsum('iiii,ij,ik,il,im->jklm', X, Y, Y, Y, Y)
def einsum_two(X, Y):
return np.einsum('iiii,ij,ik,il,im->jklm', X, Y, Y, Y, Y, optimize='optimal')
def fast_einsum(X, Y):
Z = np.einsum('ij,iiii->ij', Y, X)
W = np.einsum('il,ik->ikl', Y, Y)
Z = np.einsum('ij,im->ijm', Z, Y)
Z = np.einsum('ijm,ikl->jklm', Z, W)
return Z
Я получил fast_einsum, увидев вывод np.einsum_path('iiii,ij,ik,il,im->jklm', X, Y, Y, Y, Y, optimize='optimal')[1], который был:
Complete contraction: iiii,ij,ik,il,im->jklm
Naive scaling: 5
Optimized scaling: 5
Naive FLOP count: 1.638e+05
Optimized FLOP count: 6.662e+04
Theoretical speedup: 2.459
Largest intermediate: 4.096e+03 elements
--------------------------------------------------------------------------
scaling current remaining
--------------------------------------------------------------------------
2 ij,iiii->ij ik,il,im,ij->jklm
3 il,ik->ikl im,ij,ikl->jklm
3 ij,im->ijm ikl,ijm->jklm
5 ijm,ikl->jklm jklm->jklm
Итак, мое ожидание таково: для больших Nfast_einsum должно быть самым быстрым, а einsum_one — самым медленным. Для einsum_two должны быть некоторые накладные расходы, заметные для маленьких N, но незначительные для больших N. Однако, когда я сравнил приведенное выше, я получил следующие значения времени (числовые значения в микросекундах):
N | einsum_one | einsum_two | fast_einsum
4 15.1 949 13.8
8 256 966 87.9
12 1920 1100 683
16 7460 1180 2490
20 21800 1390 7080
Так что моя реализация fast_einsum не оптимальна. Итак, мой вопрос: как мне интерпретировать вывод np.einsum_path, чтобы я мог реализовать np.einsum как можно быстрее, не включая накладные расходы параметра optimize='optimal'?
Ответы 1
С небольшим N=10 (моя машина недостаточно велика, чтобы использовать 100 (массивы 8 МБ)
Но вот несколько быстрых таймингов:
Предварительно вычислить путь:
In [183]: path=np.einsum_path('iiii,ij,ik,il,im->jklm', X, Y, Y, Y, Y,optimize=True)[0]
по умолчанию:
In [184]: timeit M=np.einsum('iiii,ij,ik,il,im->jklm', X, Y, Y, Y, Y,optimize=False)
2.12 ms ± 4.3 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
True:
In [185]: timeit M=np.einsum('iiii,ij,ik,il,im->jklm', X, Y, Y, Y, Y,optimize=True)
509 µs ± 687 ns per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
In [186]: timeit M=np.einsum('iiii,ij,ik,il,im->jklm', X, Y, Y, Y, Y,optimize='optimal')
3.31 ms ± 25.5 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
path кажется одинаковым для True и «оптимального», но времена совершенно разные. Я не знаю почему.
Используя предварительно вычисленный путь:
In [187]: timeit M=np.einsum('iiii,ij,ik,il,im->jklm', X, Y, Y, Y, Y,optimize=path)
291 µs ± 1.07 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
In [188]: path
Out[188]: ['einsum_path', (0, 1), (0, 1), (0, 1), (0, 1)]
Поскольку все размеры одинаковы, я не уверен, что путь имеет большое значение. То есть все Y одинаковые, поэтому можно заказывать как угодно. В любом случае, расчет path в той или иной степени носит юридический характер, и не гарантируется, что он будет оптимальным для всех размеров. Как вы обнаружили, существует важный фактор масштабирования, частично из-за стоимости вычисления пути, частично (я подозреваю) из-за размера промежуточных массивов и возникающих проблем с управлением памятью.
einsum_path объясняет альтернативы:
* if a list is given that starts with ``einsum_path``, uses this as the
contraction path
* if False no optimization is taken
* if True defaults to the 'greedy' algorithm
* 'optimal' An algorithm that combinatorially explores all possible
ways of contracting the listed tensors and choosest the least costly
path. Scales exponentially with the number of terms in the
contraction.
* 'greedy' An algorithm that chooses the best pair contraction
at each step. Effectively, this algorithm searches the largest inner,
Hadamard, and then outer products at each step. Scales cubically with
the number of terms in the contraction. Equivalent to the 'optimal'
path for most contractions.
Хорошо, это объясняет, почему «оптимальный» работает медленнее — он ищет больше альтернатив. Учитывая эти размеры True, «жадный» и «оптимальный» в конечном итоге идут по одному и тому же пути.
Кажется, что оптимизация = путь — это способ увеличить скорость в einsum по сравнению с реализацией оптимизированного пути вручную, как я пытался.