Понимание NumPy einsum
Я пытаюсь понять, как работает einsum
. Я просмотрел документацию и несколько примеров, но, похоже, это не так.
Вот пример, который мы рассмотрели в классе:
C = np.einsum("ij,jk->ki", A, B)
для двух массивов A
и B
Я думаю, что это займет A^T * B
, но я не уверен (он берет транспонирование одного из них правильно?). Может ли кто-нибудь мне пройти через то, что происходит здесь (и вообще при использовании einsum
)?
Ответы
Ответ 1
(Примечание: этот ответ основан на коротком сообщении о einsum
, который я написал некоторое время назад.)
Что делает einsum
?
Представьте, что у нас есть два многомерных массива, A
и B
. Теперь предположим, что мы хотим...
- умножить
A
на B
определенным способом для создания нового массива продуктов; а затем, возможно,
- суммируйте этот новый массив вдоль отдельных осей; а затем, возможно,
- транспонировать оси нового массива в определенном порядке.
Хорошая вероятность того, что einsum
поможет нам сделать это быстрее и эффективнее с точки зрения памяти, что сочетания функций NumPy, таких как multiply
, sum
и transpose
, позволят.
Как работает einsum
?
Вот простой (но не совсем тривиальный) пример. Возьмите следующие два массива:
A = np.array([0, 1, 2])
B = np.array([[ 0, 1, 2, 3],
[ 4, 5, 6, 7],
[ 8, 9, 10, 11]])
Мы будем умножать A
и B
по элементам, а затем суммируем вдоль строк нового массива. В "нормальном" NumPy мы будем писать:
>>> (A[:, np.newaxis] * B).sum(axis=1)
array([ 0, 22, 76])
Таким образом, операция индексирования на A
выравнивает первые оси двух массивов, так что умножение может быть передано. Затем строки массива продуктов суммируются, чтобы вернуть ответ.
Теперь, если бы мы хотели использовать einsum
, мы могли бы написать:
>>> np.einsum('i,ij->i', A, B)
array([ 0, 22, 76])
Строка подписи 'i,ij->i'
является ключевым здесь и нуждается в небольшом количестве объяснений. Вы можете думать об этом в двух половинах. В левой части (слева от ->
) мы обозначили два входных массива. Справа от ->
мы обозначили массив, в который мы хотим закончить.
Вот что будет дальше:
-
A
имеет одну ось; мы отметили его i
. И B
имеет две оси; мы обозначили ось 0 как i
и ось 1 как j
.
-
повторяя метку i
в обоих массивах ввода, мы говорим einsum
, что эти две оси должны быть умножены вместе. Другими словами, мы умножаем массив A
на каждый столбец массива B
, как и на A[:, np.newaxis] * B
.
-
Обратите внимание, что j
не отображается как метка в нашем желаемом выходе; мы только что использовали i
(мы хотим получить массив 1D). Под пропуская метку, мы говорим einsum
sum вдоль этой оси. Другими словами, мы суммируем строки продуктов, как это делает .sum(axis=1)
.
Это в основном все, что вам нужно знать, чтобы использовать einsum
. Это помогает немного поиграть; если мы оставляем обе метки на выходе, 'i,ij->ij'
, мы возвращаем 2D-массив продуктов (такой же, как A[:, np.newaxis] * B
). Если мы не будем указывать выходные метки, 'i,ij->
, мы возвращаем одно число (то же самое, что и (A[:, np.newaxis] * B).sum()
).
Однако замечательная вещь в einsum
заключается в том, что сначала не создается временный массив продуктов; он просто суммирует продукты по мере их поступления. Это может привести к значительной экономии памяти.
Немного больший пример
Чтобы объяснить точечный продукт, вот два новых массива:
A = array([[1, 1, 1],
[2, 2, 2],
[5, 5, 5]])
B = array([[0, 1, 0],
[1, 1, 0],
[1, 1, 1]])
Мы вычислим произведение точек с помощью np.einsum('ij,jk->ik', A, B)
. Здесь показано изображение с меткой A
и B
и выходным массивом, который мы получаем от функции:
![введите описание изображения здесь]()
Вы можете видеть, что метка j
повторяется - это означает, что мы умножаем строки A
на столбцы B
. Кроме того, метка j
не включается в вывод - поэтому мы суммируем эти продукты. Ярлыки i
и k
сохраняются для вывода, поэтому мы возвращаем 2D-массив.
Возможно, было бы еще проще сравнить этот результат с массивом, где метка j
не суммируется. Ниже, слева, вы можете увидеть 3D-массив, полученный в результате записи np.einsum('ij,jk->ijk', A, B)
(т.е. Мы сохранили метку j
):
![введите описание изображения здесь]()
Суммирование оси j
дает ожидаемый точечный продукт, показанный справа.
Некоторые упражнения
Чтобы узнать больше о einsum
, полезно реализовать знакомые операции с массивом NumPy, используя нотацию индекса. Все, что связано с комбинациями умножающих и суммирующих осей, можно записать с помощью einsum
.
Пусть A и B - два 1D массива с одинаковой длиной. Например, A = np.arange(10)
и B = np.arange(5, 15)
.
-
Сумма A
может быть записана:
np.einsum('i->', A)
-
Элементарное умножение, A * B
, можно записать:
np.einsum('i,i->i', A, B)
-
Внутренний продукт или точечный продукт np.inner(A, B)
или np.dot(A, B)
можно записать:
np.einsum('i,i->', A, B) # or just use 'i,i'
-
Внешнее произведение np.outer(A, B)
можно записать:
np.einsum('i,j->ij', A, B)
Для 2D-массивов C
и D
при условии, что оси являются совместимыми длинами (как одна и та же длина, либо одна из них имеет длину 1), вот несколько примеров:
-
След C
(сумма основной диагонали), np.trace(C)
, можно записать:
np.einsum('ii', C)
-
Элементное умножение C
и транспонирование D
, C * D.T
можно записать:
np.einsum('ij,ji->ij', C, D)
-
Умножение каждого элемента C
на массив D
(для создания 4D-массива), C[:, :, None, None] * D
, можно записать:
np.einsum('ij,kl->ijkl', C, D)
Ответ 2
Позволяет сделать 2 массива с разными, но совместимыми размерами, чтобы выделить их взаимодействие
In [43]: A=np.arange(6).reshape(2,3)
Out[43]:
array([[0, 1, 2],
[3, 4, 5]])
In [44]: B=np.arange(12).reshape(3,4)
Out[44]:
array([[ 0, 1, 2, 3],
[ 4, 5, 6, 7],
[ 8, 9, 10, 11]])
Ваш расчет принимает "точку" (сумму произведений) (2,3) с (3,4) для создания массива (4,2). i
- 1-й тусклый A
, последний из C
; k
последний из B
, 1-й из C
. j
"поглощается" суммированием.
In [45]: C=np.einsum('ij,jk->ki',A,B)
Out[45]:
array([[20, 56],
[23, 68],
[26, 80],
[29, 92]])
Это то же самое, что и np.dot(A,B).T
- это конечный вывод, который был транспонирован.
Чтобы узнать больше о том, что происходит с j
, измените индексы C
на ijk
:
In [46]: np.einsum('ij,jk->ijk',A,B)
Out[46]:
array([[[ 0, 0, 0, 0],
[ 4, 5, 6, 7],
[16, 18, 20, 22]],
[[ 0, 3, 6, 9],
[16, 20, 24, 28],
[40, 45, 50, 55]]])
Это также можно создать с помощью
A[:,:,None]*B[None,:,:]
То есть добавьте размер k
в конец A
и i
в начало B
, в результате получим массив (2,3,4).
0 + 4 + 16 = 20
, 9 + 28 + 55 = 92
и т.д.; Сумма на j
и транспонирование для получения более раннего результата:
np.sum(A[:,:,None] * B[None,:,:], axis=1).T
# C[k,i] = sum(j) A[i,j (,k) ] * B[(i,) j,k]
Ответ 3
Я нашел NumPy: трюки торговли (часть II) поучительный
Мы используем → для указания порядка выходного массива. Поэтому подумайте, что ij, i- > j 'имеют левую сторону (LHS) и правую сторону (RHS). Любое повторение меток на LHS вычисляет элемент продукта мудрый, а затем суммируется. Изменив метку на стороне RHS (выход), мы можем определить ось, в которой мы хотим действовать относительно входного массива, т.е. Суммирование вдоль оси 0, 1 и т.д.
import numpy as np
>>> a
array([[1, 1, 1],
[2, 2, 2],
[3, 3, 3]])
>>> b
array([[0, 1, 2],
[3, 4, 5],
[6, 7, 8]])
>>> d = np.einsum('ij, jk->ki', a, b)
Обратите внимание, что существует три оси, i, j, k, а j повторяется (слева). i,j
представляют строки и столбцы для a
. j,k
для b
.
Чтобы вычислить произведение и выровнять ось j
, нам нужно добавить ось к a
. (b
будет передаваться вдоль (?) первой оси)
a[i, j, k]
b[j, k]
>>> c = a[:,:,np.newaxis] * b
>>> c
array([[[ 0, 1, 2],
[ 3, 4, 5],
[ 6, 7, 8]],
[[ 0, 2, 4],
[ 6, 8, 10],
[12, 14, 16]],
[[ 0, 3, 6],
[ 9, 12, 15],
[18, 21, 24]]])
j
отсутствует в правой части, поэтому мы суммируем по j
, которая является второй осью массива 3x3x3
>>> c = c.sum(1)
>>> c
array([[ 9, 12, 15],
[18, 24, 30],
[27, 36, 45]])
Наконец, индексы (в алфавитном порядке) меняются на правую сторону, поэтому мы транспонируем.
>>> c.T
array([[ 9, 18, 27],
[12, 24, 36],
[15, 30, 45]])
>>> np.einsum('ij, jk->ki', a, b)
array([[ 9, 18, 27],
[12, 24, 36],
[15, 30, 45]])
>>>
Ответ 4
Вот несколько примеров, иллюстрирующих использование np.einsum()
в реализации некоторых общих тензорных или n-мерных операций.
Входы
In [197]: vec
Out[197]: array([0, 1, 2, 3])
In [198]: A
Out[198]:
array([[11, 12, 13, 14],
[21, 22, 23, 24],
[31, 32, 33, 34],
[41, 42, 43, 44]])
In [199]: B
Out[199]:
array([[1, 1, 1, 1],
[2, 2, 2, 2],
[3, 3, 3, 3],
[4, 4, 4, 4]])
1) матричное умножение (похожее на np.matmul(arr1, arr2)
)
In [200]: np.einsum("ij, jk -> ik", A, B)
Out[200]:
array([[130, 130, 130, 130],
[230, 230, 230, 230],
[330, 330, 330, 330],
[430, 430, 430, 430]])
2) Извлеките элементы вдоль основной диагонали (аналогично np.diag(arr)
)
In [202]: np.einsum("ii -> i", A)
Out[202]: array([11, 22, 33, 44])
3) Продукт Адамара (т.е. элементное произведение двух массивов) (похоже на arr1 * arr2
)
In [203]: np.einsum("ij, ij -> ij", A, B)
Out[203]:
array([[ 11, 12, 13, 14],
[ 42, 44, 46, 48],
[ 93, 96, 99, 102],
[164, 168, 172, 176]])
4) Элементарное возведение в квадрат (похожее на np.square(arr)
или arr ** 2
)
In [210]: np.einsum("ij, ij -> ij", B, B)
Out[210]:
array([[ 1, 1, 1, 1],
[ 4, 4, 4, 4],
[ 9, 9, 9, 9],
[16, 16, 16, 16]])
5) Трассировка (т.е. сумма основных диагональных элементов) (аналогично np.trace(arr)
)
In [217]: np.einsum("ii -> ", A)
Out[217]: 110
6) Матрица транспонирует (аналогично np.transpose(arr)
)
In [221]: np.einsum("ij -> ji", A)
Out[221]:
array([[11, 21, 31, 41],
[12, 22, 32, 42],
[13, 23, 33, 43],
[14, 24, 34, 44]])
7) Внешний продукт (векторов) (аналогично np.outer(vec1, vec2)
)
In [255]: np.einsum("i, j -> ij", vec, vec)
Out[255]:
array([[0, 0, 0, 0],
[0, 1, 2, 3],
[0, 2, 4, 6],
[0, 3, 6, 9]])
8) Внутренний продукт (векторов) (аналогично np.inner(vec1, vec2)
)
In [256]: np.einsum("i, i -> ", vec, vec)
Out[256]: 14
9) Сумма вдоль оси 0 (похожа на np.sum(arr, axis=0)
)
In [260]: np.einsum("ij -> j", B)
Out[260]: array([10, 10, 10, 10])
10) Сумма вдоль оси 1 (аналогично np.sum(arr, axis=1)
)
In [261]: np.einsum("ij -> i", B)
Out[261]: array([ 4, 8, 12, 16])
11) Массовое умножение матрицы
In [287]: BM = np.stack((A, B), axis=0)
In [288]: BM
Out[288]:
array([[[11, 12, 13, 14],
[21, 22, 23, 24],
[31, 32, 33, 34],
[41, 42, 43, 44]],
[[ 1, 1, 1, 1],
[ 2, 2, 2, 2],
[ 3, 3, 3, 3],
[ 4, 4, 4, 4]]])
In [289]: BM.shape
Out[289]: (2, 4, 4)
# batch matrix multiply using einsum
In [292]: BMM = np.einsum("bij, bjk -> bik", BM, BM)
In [293]: BMM
Out[293]:
array([[[1350, 1400, 1450, 1500],
[2390, 2480, 2570, 2660],
[3430, 3560, 3690, 3820],
[4470, 4640, 4810, 4980]],
[[ 10, 10, 10, 10],
[ 20, 20, 20, 20],
[ 30, 30, 30, 30],
[ 40, 40, 40, 40]]])
In [294]: BMM.shape
Out[294]: (2, 4, 4)
12) суммировать вдоль оси 2 (аналогично np.sum(arr, axis=2)
)
In [330]: np.einsum("ijk -> ij", BM)
Out[330]:
array([[ 50, 90, 130, 170],
[ 4, 8, 12, 16]])
13) суммировать все элементы в массиве (аналогично np.sum(arr)
)
In [335]: np.einsum("ijk -> ", BM)
Out[335]: 480
14) сумма по нескольким осям (т.е. маргинализация)
(аналогично np.sum(arr, axis=(axis0, axis1, axis2, axis3, axis4, axis6, axis7))
)
# 8D array
In [354]: R = np.random.standard_normal((3,5,4,6,8,2,7,9))
# marginalize out axis 5 (i.e. "n" here)
In [363]: esum = np.einsum("ijklmnop -> n", R)
# marginalize out axis 5 (i.e. sum over rest of the axes)
In [364]: nsum = np.sum(R, axis=(0,1,2,3,4,6,7))
In [365]: np.allclose(esum, nsum)
Out[365]: True
15) Double Dot Products (аналогично np.sum(hadamard-product) cf. 3)
In [772]: A
Out[772]:
array([[1, 2, 3],
[4, 2, 2],
[2, 3, 4]])
In [773]: B
Out[773]:
array([[1, 4, 7],
[2, 5, 8],
[3, 6, 9]])
In [774]: np.einsum("ij, ij -> ", A, B)
Out[774]: 124
Подробнее здесь: Эйнштейн-Суммирование и определенно здесь: Тензорная-нотация