Capire l'einsum di NumPy

Sto lottando per capire esattamente come einsumfunziona. Ho esaminato la documentazione e alcuni esempi, ma non sembra aderire.

Ecco un esempio che abbiamo esaminato in classe:

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

per due array AeB

Penso che questo richiederebbe A^T * B, ma non sono sicuro (sta prendendo il recepimento di uno di loro giusto?). Qualcuno può guidarmi attraverso ciò che sta accadendo qui (e in generale durante l'utilizzo einsum)?

(Nota: questa risposta si basa su un breve post sul blog cheeinsum ho scritto qualche tempo fa.)

Cosa fa einsum?

Immagina di avere due array multidimensionali Ae B. Ora supponiamo di voler ...

• moltiplicarsi A con Bin un modo particolare per creare una nuova gamma di prodotti; e poi forse
• sommare questo nuovo array lungo determinati assi; e poi forse
• trasporre gli assi del nuovo array in un ordine particolare.

C'è una buona probabilità che einsumci aiuterà a fare questo più veloce e più memoria-efficiente che le combinazioni delle funzioni NumPy piace multiply, sume transposeconsentirà.

Come einsumfunziona?

Ecco un esempio semplice (ma non del tutto banale). Prendi i seguenti due array:

A = np.array([0, 1, 2])

B = np.array([[ 0,  1,  2,  3],
              [ 4,  5,  6,  7],
              [ 8,  9, 10, 11]])

Moltiplicheremo Ae saggieremo gli Belementi e poi la somma lungo le righe del nuovo array. In "normale" NumPy scriveremmo:

>>> (A[:, np.newaxis] * B).sum(axis=1)
array([ 0, 22, 76])

Quindi qui l'operazione di indicizzazione Aallinea i primi assi dei due array in modo che la moltiplicazione possa essere trasmessa. Le righe della matrice di prodotti vengono quindi sommate per restituire la risposta.

Ora, se volessimo usare einsuminvece, potremmo scrivere:

>>> np.einsum('i,ij->i', A, B)
array([ 0, 22, 76])

La stringa della firma'i,ij->i' è la chiave qui e ha bisogno di un po 'di spiegazione. Puoi pensarci in due metà. Sul lato sinistro (a sinistra del ->) abbiamo etichettato i due array di input. A destra ->, abbiamo etichettato l'array con cui vogliamo finire.

Ecco cosa succede dopo:

• Aha un asse; l'abbiamo etichettato i. E Bha due assi; abbiamo etichettato l'asse 0 come ie l'asse 1 come j.

• Da ripetere l'etichetta iin entrambi gli array di input, stiamo dicendo einsumche questi due assi devono essere moltiplicati insieme. In altre parole, stiamo moltiplicando l'array Aper ogni colonna dell'array B, proprio come A[:, np.newaxis] * Bfa.

• Si noti che jnon appare come un'etichetta nell'output desiderato; abbiamo appena usato i(vogliamo finire con un array 1D). Con omettendo l'etichetta, stiamo dicendo einsuma sommare lungo questo asse. In altre parole, stiamo sommando le file dei prodotti, proprio come .sum(axis=1)fa.

Questo è fondamentalmente tutto ciò che devi sapere per usare einsum. Aiuta a giocare un po '; se lasciamo entrambe le etichette nell'output, 'i,ij->ij'otteniamo una matrice di prodotti 2D (uguale a A[:, np.newaxis] * B). Se diciamo che non ci sono etichette di output, 'i,ij->otteniamo un singolo numero (come nel fare (A[:, np.newaxis] * B).sum()).

La cosa grandiosa einsumè che non crea prima una serie temporanea di prodotti; somma semplicemente i prodotti così come sono. Questo può portare a grandi risparmi nell'uso della memoria.

Un esempio leggermente più grande

Per spiegare il prodotto punto, ecco due nuovi array:

A = array([[1, 1, 1],
           [2, 2, 2],
           [5, 5, 5]])

B = array([[0, 1, 0],
           [1, 1, 0],
           [1, 1, 1]])

Calcoleremo il prodotto dot usando np.einsum('ij,jk->ik', A, B). Ecco una foto che mostra l'etichettatura del Ae Be l'array di output che si ottiene dalla funzione:

Puoi vedere che l'etichetta jviene ripetuta - questo significa che stiamo moltiplicando le righe di Acon le colonne di B. Inoltre, l'etichetta jnon è inclusa nell'output: stiamo sommando questi prodotti. Le etichette ie kvengono conservate per l'output, quindi torniamo a un array 2D.

Potrebbe essere ancora più chiaro confrontare questo risultato con l'array in cui l'etichetta nonj viene sommata. Sotto, a sinistra, puoi vedere l'array 3D che risulta dalla scrittura (ovvero abbiamo mantenuto l'etichetta ):np.einsum('ij,jk->ijk', A, B)j

L'asse di somma jfornisce il prodotto punto atteso, mostrato a destra.

Alcuni esercizi

Per avere più senso einsum, può essere utile implementare le operazioni familiari dell'array NumPy usando la notazione in pedice. Tutto ciò che coinvolge combinazioni di moltiplicare e sommare gli assi può essere scritto usando einsum.

Sia A e B due array 1D con la stessa lunghezza. Ad esempio, A = np.arange(10)e B = np.arange(5, 15).

• La somma di Apuò essere scritta:

  np.einsum('i->', A)

• La moltiplicazione degli elementi A * B, può essere scritta:

  np.einsum('i,i->i', A, B)

• Il prodotto interno o il prodotto punto, np.inner(A, B)oppure np.dot(A, B), può essere scritto:

  np.einsum('i,i->', A, B) # or just use 'i,i'

• Il prodotto esterno np.outer(A, B), può essere scritto:

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

Per gli array 2D Ce D, a condizione che gli assi siano lunghezze compatibili (sia la stessa lunghezza che una di esse ha lunghezza 1), ecco alcuni esempi:

• La traccia di C(somma della diagonale principale) np.trace(C), può essere scritta:

  np.einsum('ii', C)

• Element-saggio moltiplicazione Ce la trasposta di D, C * D.T, può essere scritta:

  np.einsum('ij,ji->ij', C, D)

• Moltiplicando ciascun elemento di Cper l'array D(per creare un array 4D) C[:, :, None, None] * D, è possibile scrivere:

  np.einsum('ij,kl->ijkl', C, D)  

Afferrare l'idea di numpy.einsum()è molto facile se la capisci in modo intuitivo. A titolo di esempio, iniziamo con una semplice descrizione che coinvolge la moltiplicazione di matrici .

Per usare numpy.einsum(), tutto ciò che devi fare è passare la cosiddetta stringa dei pedici come argomento, seguita dai tuoi array di input .

Diciamo che avete due 2D array, Ae B, e si vuole fare la moltiplicazione di matrici. Quindi fai:

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

Qui la stringa del pediceij corrisponde alla matrice Amentre la stringa del pedicejk corrisponde alla matrice B. Inoltre, la cosa più importante da notare qui è che il numero di caratteri in ciascuna stringa di pedice deve corrispondere alle dimensioni dell'array. (ovvero due caratteri per array 2D, tre caratteri per array 3D e così via). E se si ripetono i caratteri tra stringhe di indice ( jnel nostro caso), ciò significa che si desidera che la einsomma avvenga lungo tali dimensioni. Pertanto, saranno ridotti in somma. (cioè quella dimensione sparirà )

La stringa del pedice dopo questo ->, sarà la nostra matrice risultante. Se lo lasci vuoto, tutto verrà sommato e di conseguenza verrà restituito un valore scalare. Altrimenti l'array risultante avrà dimensioni in base alla stringa del pedice . Nel nostro esempio, lo sarà ik. Questo è intuitivo perché sappiamo che per la moltiplicazione delle matrici il numero di colonne nell'array Adeve corrispondere al numero di righe nell'array Bche è ciò che sta accadendo qui (ovvero codifichiamo questa conoscenza ripetendo il carattere jnella stringa del pedice )

Ecco alcuni altri esempi che illustrano in modo succinto l'uso / la potenza np.einsum()dell'implementazione di alcune comuni operazioni tensore o nd-array .

ingressi

# 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) Moltiplicazione della matrice (simile a 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) Estrai gli elementi lungo la diagonale principale (simile a np.diag(arr))

In [202]: np.einsum("ii -> i", A)
Out[202]: array([11, 22, 33, 44])

3) Prodotto Hadamard (ovvero un prodotto saggio di due array) (simile a 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) Quadratura a livello di elemento (simile a np.square(arr)o 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) Traccia (ovvero somma degli elementi diagonali principali) (simile a np.trace(arr))

In [217]: np.einsum("ii -> ", A)
Out[217]: 110

6) Trasposizione matrice (simile a 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) Prodotto esterno (di vettori) (simile a 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) Prodotto interno (di vettori) (simile a np.inner(vec1, vec2))

In [256]: np.einsum("i, i -> ", vec, vec)
Out[256]: 14

9) Somma lungo l'asse 0 (simile a np.sum(arr, axis=0))

In [260]: np.einsum("ij -> j", B)
Out[260]: array([10, 10, 10, 10])

10) Somma lungo l'asse 1 (simile a np.sum(arr, axis=1))

In [261]: np.einsum("ij -> i", B)
Out[261]: array([ 4,  8, 12, 16])

11) Moltiplicazione della matrice batch

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) Somma lungo l'asse 2 (simile a np.sum(arr, axis=2))

In [330]: np.einsum("ijk -> ij", BM)
Out[330]: 
array([[ 50,  90, 130, 170],
       [  4,   8,  12,  16]])

13) Somma tutti gli elementi nella matrice (simile a np.sum(arr))

In [335]: np.einsum("ijk -> ", BM)
Out[335]: 480

14) Somma su più assi (es. Emarginazione)
(simile a 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) Prodotti a doppio punto (simile a np.sum (prodotto hadamard) cfr. 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) Moltiplicazione di array 2D e 3D

Tale moltiplicazione potrebbe essere molto utile quando si risolve il sistema lineare di equazioni ( Ax = b ) in cui si desidera verificare il risultato.

# 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

Al contrario, se si deve usare np.matmul()per questa verifica, dobbiamo fare un paio di reshapeoperazioni per ottenere lo stesso risultato come:

# 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

Bonus : leggi ulteriori informazioni matematiche qui: Einstein-Summation e sicuramente qui: Tensor-Notation

Consente di creare 2 array, con dimensioni diverse ma compatibili per evidenziare la loro interazione

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]])

Il tuo calcolo richiede un "punto" (somma dei prodotti) di un (2,3) con un (3,4) per produrre un array (4,2). iè il 1o fioco di A, l'ultimo di C; kl'ultimo di B, il 1o di C. jviene "consumato" dalla somma.

In [45]: C=np.einsum('ij,jk->ki',A,B)
Out[45]: 
array([[20, 56],
       [23, 68],
       [26, 80],
       [29, 92]])

È lo stesso di np.dot(A,B).T: è l'output finale che viene trasposto.

Per vedere di più di ciò che accade j, modifica gli Cindici in 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]]])

Questo può anche essere prodotto con:

A[:,:,None]*B[None,:,:]

Cioè, aggiungi una kdimensione alla fine di Ae una iin primo piano B, risultando in una matrice (2,3,4).

0 + 4 + 16 = 20, 9 + 28 + 55 = 92ecc .; Somma je trasponi per ottenere il risultato precedente:

np.sum(A[:,:,None] * B[None,:,:], axis=1).T

# C[k,i] = sum(j) A[i,j (,k) ] * B[(i,)  j,k]

Ho trovato NumPy: i trucchi del mestiere (parte II) istruttivi

Usiamo -> per indicare l'ordine dell'array di output. Quindi pensa a 'ij, i-> j' come se avesse il lato sinistro (LHS) e il lato destro (RHS). Qualsiasi ripetizione di etichette sull'LHS calcola l'elemento del prodotto in modo saggio e quindi si somma. Modificando l'etichetta sul lato RHS (output), possiamo definire l'asse in cui vogliamo procedere rispetto all'array di input, ovvero la somma lungo l'asse 0, 1 e così via.

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)

Notare che ci sono tre assi, i, j, k, e che j viene ripetuto (sul lato sinistro). i,jrappresentano righe e colonne per a. j,kper b.

Per calcolare il prodotto e allineare l' jasse, è necessario aggiungere un asse a. ( bverrà trasmesso lungo (?) il primo asse)

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è assente dal lato destro, quindi si somma su jquale sia il secondo asse dell'array 3x3x3

>>> c = c.sum(1)
>>> c
array([[ 9, 12, 15],
       [18, 24, 30],
       [27, 36, 45]])

Infine, gli indici sono (in ordine alfabetico) invertiti sul lato destro, quindi trasponiamo.

>>> 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]])
>>>

Quando ho letto le equazioni di einsum, ho trovato più utile essere in grado di ridurle mentalmente alle loro versioni imperative.

Cominciamo con la seguente (imponente) istruzione:

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

Lavorando prima attraverso la punteggiatura vediamo che abbiamo due BLOB separati da virgola di 4 lettere - bhwie bhwj, prima della freccia, e un singolo BLOB di 3 lettere bijdopo di esso. Pertanto, l'equazione produce un risultato tensore di grado 3 da due ingressi tensore di grado 4.

Ora, lascia che ogni lettera in ogni BLOB sia il nome di una variabile di intervallo. La posizione in cui la lettera appare nel BLOB è l'indice dell'asse su cui si estende in quel tensore. La somma imperativa che produce ogni elemento di C, quindi, deve iniziare con tre cicli nidificati, uno per ogni indice di C.

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] = ...

Quindi, essenzialmente, hai un forciclo per ogni indice di output di C. Per ora lasceremo gli intervalli indeterminati.

Ora guardiamo al lato sinistro - ci sono delle variabili di intervallo lì che non compaiono sul lato destro ? Nel nostro caso - sì, he w. Aggiungi un forciclo annidato interno per ciascuna di tali variabili:

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(...):
                    ...

All'interno del ciclo più interno ora abbiamo tutti gli indici definiti, quindi possiamo scrivere la somma effettiva e la traduzione è completa:

# 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]

Se finora sei stato in grado di seguire il codice, allora congratulazioni! Questo è tutto ciò che serve per poter leggere le equazioni di einsum. Si noti in particolare come la formula originale di einsum sia mappata alla dichiarazione di sommatoria finale nel frammento sopra. I for-loop e i limiti del range sono solo soffici e quell'ultima affermazione è tutto ciò di cui hai veramente bisogno per capire cosa sta succedendo.

Per completezza, vediamo come determinare gli intervalli per ciascuna variabile di intervallo. Bene, l'intervallo di ogni variabile è semplicemente la lunghezza delle dimensioni che indicizza. Ovviamente, se una variabile indicizza più di una dimensione in uno o più tensori, allora le lunghezze di ciascuna di quelle dimensioni devono essere uguali. Ecco il codice sopra con gli intervalli completi:

# 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]

Penso che l'esempio più semplice sia nei documenti di Tensorflow

Esistono quattro passaggi per convertire l'equazione in notazione di einsum. Prendiamo questa equazione come esempioC[i,k] = sum_j A[i,j] * B[j,k]

1. Per prima cosa eliminiamo i nomi delle variabili. Noi abbiamoik = sum_j ij * jk
2. Lasciamo cadere il sum_jtermine perché è implicito. Noi abbiamoik = ij * jk
3. Sostituiamo *con ,. Noi abbiamoik = ij, jk
4. L'uscita è su RHS ed è separata da ->segno. Noi abbiamoij, jk -> ik

L'interprete einsum esegue semplicemente questi 4 passaggi al contrario. Tutti gli indici mancanti nel risultato vengono sommati.

Ecco alcuni altri esempi dai documenti

# 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]