Розуміння ейнсума NumPy

Я намагаюся зрозуміти, як саме einsumпрацює. Я переглянув документацію та кілька прикладів, але це, здається, не дотримується.

Ось приклад, який ми перейшли на уроці:

C = np.einsum("ij,jk->ki", A, B)

для двох масивів AіB

Я думаю, що це зайняло б A^T * B, але я не впевнений (це перенесення одного з них правильно?). Чи може хтось переглядати мене саме через те, що відбувається тут (і взагалі при використанні einsum)?

(Примітка. Ця відповідь заснована на короткій публікації в блозі, про яку 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(axis=1)і.

Це в основному все, що потрібно знати, щоб використовувати einsum. Це допомагає трохи пограти; якщо ми залишимо обидві мітки у висновку, 'i,ij->ij'ми отримаємо назад двовимірний масив продуктів (такий же, як 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знаходиться НЕ підбивалися. Нижче зліва ви бачите тривимірний масив, який є результатом написання 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)

Для двовимірних масивів 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)  

Осягнути ідею numpy.einsum()дуже просто, якщо зрозуміти її інтуїтивно. Як приклад, почнемо з простого опису, що включає множення матриць .

Для використання numpy.einsum()все, що вам потрібно зробити, - це передавати так званий рядок підписок як аргумент, за яким слід вводити масиви введення .

Скажімо, у вас є два 2D масиви, Aі B, і ви хочете зробити матричне множення. Отже, ви робите:

np.einsum("ij, jk -> ik", A, B)

Тут рядок підрядного індексу ij відповідає масиву, в Aтой час як рядок підрядного індексу jk відповідає масиву B. Також найважливіше, що тут слід зазначити, - це те, що кількість символів у кожному рядку індекса повинна відповідати розмірам масиву. (Тобто два символу для 2D масивів, три CHARS для 3D масивів, і так далі.) І якщо ви повторите символи між індексними рядками ( jв нашому випадку), то це означає , що ви хочете, щоб einсума відбудуться уздовж цих розмірів. Таким чином, вони будуть зменшені на суму. (Тобто розмір буде пішов )

Підрядковий рядок після цього ->, буде наш результуючий масив. Якщо ви залишите його порожнім, то все підсумовується і в результаті повертається скалярне значення. В іншому випадку масив, що отримує результат, матиме розміри відповідно до рядка підпису . У нашому прикладі це буде ik. Це інтуїтивно, тому що ми знаємо, що для множення матриці кількість стовпців у масиві Aмає відповідати кількості рядків у масиві B, що саме тут відбувається (тобто ми кодуємо це знання повторенням символу jв рядку підрядника )

Ось ще кілька прикладів, що ілюструють використання / потужність для np.einsum()здійснення деяких загальних тензорних або nd-масивів операцій, стисло.

Вхідні дані

# a vector
In [197]: vec
Out[197]: array([0, 1, 2, 3])

# an array
In [198]: A
Out[198]: 
array([[11, 12, 13, 14],
       [21, 22, 23, 24],
       [31, 32, 33, 34],
       [41, 42, 43, 44]])

# another array
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) Продукти з подвійною крапкою (схожі на np.sum (hadamard-product), пор. 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

16) 2D та 3D множення масивів

Таке множення може бути дуже корисним при вирішенні лінійної системи рівнянь ( Ax = b ), де потрібно перевірити результат.

# inputs
In [115]: A = np.random.rand(3,3)
In [116]: b = np.random.rand(3, 4, 5)

# solve for x
In [117]: x = np.linalg.solve(A, b.reshape(b.shape[0], -1)).reshape(b.shape)

# 2D and 3D array multiplication :)
In [118]: Ax = np.einsum('ij, jkl', A, x)

# indeed the same!
In [119]: np.allclose(Ax, b)
Out[119]: True

Навпаки, якщо вам доведеться використовувати np.matmul()для цієї перевірки, ми повинні зробити пару reshapeоперацій, щоб досягти такого ж результату, як:

# reshape 3D array `x` to 2D, perform matmul
# then reshape the resultant array to 3D
In [123]: Ax_matmul = np.matmul(A, x.reshape(x.shape[0], -1)).reshape(x.shape)

# indeed correct!
In [124]: np.allclose(Ax, Ax_matmul)
Out[124]: True

Бонус : Детальніше про математику читайте тут: Ейнштейн-Підсумовування і, безумовно, тут: Тензор-Нотація

Дозволяє зробити 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останній з B1-го 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]

Я знайшов 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]])
>>>

Читаючи рівняння ейнсума, я вважаю, що найкориснішим є можливість лише подумки звести їх до своїх імперативних версій.

Почнемо із наступного (нав'язуючого) твердження:

C = np.einsum('bhwi,bhwj->bij', A, B)

Опрацьовуючи розділові знаки спочатку, ми бачимо, що у нас є дві 4- літерні крапки, розділені комами - bhwiі bhwj, перед стрілкою, і одна крапка bijпісля неї. Тому рівняння виробляє результат тензору ранжу-3 з двох входів тензору рангу-4.

Тепер, нехай кожна буква в кожному краплі є назвою змінної діапазону. Положення, в якому буква з’являється в краплі, є індексом осі, по якій вона перебуває в цьому тензорі. Таким чином, імперативне підсумовування, що створює кожен елемент С, має починатися з трьох вкладених циклів, по одній для кожного показника С.

for b in range(...):
    for i in range(...):
        for j in range(...):
            # the variables b, i and j index C in the order of their appearance in the equation
            C[b, i, j] = ...

Отже, по суті, у вас є for цикл для кожного вихідного індексу C. Ми покинемо діапазони невизначеними.

Далі ми подивимось на ліву частину - чи є там змінні діапазону, які не відображаються на правій частині? У нашому випадку - так, hі w. Додайте внутрішній вкладений forцикл для кожної такої змінної:

for b in range(...):
    for i in range(...):
        for j in range(...):
            C[b, i, j] = 0
            for h in range(...):
                for w in range(...):
                    ...

Всередині внутрішньої петлі у нас визначені всі індекси, тому ми можемо записати фактичне підсумок і переклад завершено:

# three nested for-loops that index the elements of C
for b in range(...):
    for i in range(...):
        for j in range(...):

            # prepare to sum
            C[b, i, j] = 0

            # two nested for-loops for the two indexes that don't appear on the right-hand side
            for h in range(...):
                for w in range(...):
                    # Sum! Compare the statement below with the original einsum formula
                    # 'bhwi,bhwj->bij'

                    C[b, i, j] += A[b, h, w, i] * B[b, h, w, j]

Якщо ви досі змогли дотримуватися коду, то вітайте! Це все, що вам потрібно, щоб мати можливість читати рівняння ейнсума. Зокрема, зверніть увагу на те, як оригінальна формула einsum відображається до остаточного висловлення підсумків у фрагменті вище. Межі for-циклів і межі діапазону - це просто пух, і це остаточне твердження - все, що вам потрібно, щоб зрозуміти, що відбувається.

Для повноти подивимось, як визначити діапазони для кожної змінної діапазону. Ну, діапазон кожної змінної - це просто довжина розмірності, яку вона індексує. Очевидно, якщо змінна індексує більше одного виміру в одному або декількох тензорах, то довжини кожного з цих розмірів повинні бути рівними. Ось наведений вище код із повним діапазоном:

# C's shape is determined by the shapes of the inputs
# b indexes both A and B, so its range can come from either A.shape or B.shape
# i indexes only A, so its range can only come from A.shape, the same is true for j and B
assert A.shape[0] == B.shape[0]
assert A.shape[1] == B.shape[1]
assert A.shape[2] == B.shape[2]
C = np.zeros((A.shape[0], A.shape[3], B.shape[3]))
for b in range(A.shape[0]): # b indexes both A and B, or B.shape[0], which must be the same
    for i in range(A.shape[3]):
        for j in range(B.shape[3]):
            # h and w can come from either A or B
            for h in range(A.shape[1]):
                for w in range(A.shape[2]):
                    C[b, i, j] += A[b, h, w, i] * B[b, h, w, j]

Я думаю, що найпростіший приклад - у документах з tensorflow

Існує чотири кроки для перетворення рівняння в нотацію einsum. Давайте візьмемо це рівняння як прикладC[i,k] = sum_j A[i,j] * B[j,k]

1. Спочатку ми скидаємо імена змінних. Ми отримуємоik = sum_j ij * jk
2. Ми скидаємо sum_jтермін, як він неявний. Ми отримуємоik = ij * jk
3. Заміняємо *на ,. Ми отримуємоik = ij, jk
4. Вихід знаходиться на RHS і розділений ->знаком. Ми отримуємоij, jk -> ik

Інтерпретатор einsum просто виконує ці 4 кроки в зворотному порядку. Усі індекси, відсутні в результаті, підсумовуються.

Ось ще кілька прикладів із документів

# Matrix multiplication
einsum('ij,jk->ik', m0, m1)  # output[i,k] = sum_j m0[i,j] * m1[j, k]

# Dot product
einsum('i,i->', u, v)  # output = sum_i u[i]*v[i]

# Outer product
einsum('i,j->ij', u, v)  # output[i,j] = u[i]*v[j]

# Transpose
einsum('ij->ji', m)  # output[j,i] = m[i,j]

# Trace
einsum('ii', m)  # output[j,i] = trace(m) = sum_i m[i, i]

# Batch matrix multiplication
einsum('aij,ajk->aik', s, t)  # out[a,i,k] = sum_j s[a,i,j] * t[a, j, k]