NumPys einsum verstehen

Ich habe Mühe, genau zu verstehen, wie es einsumfunktioniert. Ich habe mir die Dokumentation und einige Beispiele angesehen, aber es scheint nicht zu bleiben.

Hier ist ein Beispiel, das wir im Unterricht besprochen haben:

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

für zwei Arrays AundB

Ich denke, das würde dauern A^T * B, aber ich bin mir nicht sicher (es braucht die Transponierung von einem von ihnen, oder?). Kann mich jemand genau durch das führen, was hier passiert (und im Allgemeinen bei der Verwendung einsum)?

(Hinweis: Diese Antwort ist auf einer kurzen basierende Blog - Post über die einsumich vor einer Weile geschrieben.)

Was macht einsumdas

Stellen Sie sich vor, wir haben zwei mehrdimensionale Arrays Aund B. Nehmen wir jetzt an, wir wollen ...

• mehrfach A mit Bin einer bestimmten Art und Weise neue Reihe von Produkten zu erzeugen; und dann vielleicht
• Summe dieses neue Array entlang bestimmter Achsen; und dann vielleicht
• transponieren Sie die Achsen des neuen Arrays in einer bestimmten Reihenfolge.

Es gibt eine gute Chance , dass einsumuns helfen wird , tun dies schneller und Speicher-effizient , dass Kombinationen der NumPy Funktionen wie multiply, sumund transposeermöglicht.

Wie funktioniert das einsum?

Hier ist ein einfaches (aber nicht ganz triviales) Beispiel. Nehmen Sie die folgenden zwei Arrays:

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

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

Wir werden multiplizieren Aund Belementweise und dann entlang der Zeilen des neuen Arrays summieren. In "normalem" NumPy würden wir schreiben:

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

Hier richtet die Indizierungsoperation Adie ersten Achsen der beiden Arrays so aus, dass die Multiplikation gesendet werden kann. Die Zeilen des Produktarrays werden dann summiert, um die Antwort zurückzugeben.

Wenn wir einsumstattdessen verwenden wollten , könnten wir schreiben:

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

Der Signatur - String 'i,ij->i'ist hier der Schlüssel und braucht ein wenig zu erklären. Sie können sich das in zwei Hälften vorstellen. Auf der linken Seite (links von ->) haben wir die beiden Eingabearrays beschriftet. Rechts von ->haben wir das Array beschriftet, mit dem wir enden möchten.

Folgendes passiert als nächstes:

• Ahat eine Achse; wir haben es beschriftet i. Und Bhat zwei Achsen; Wir haben Achse 0 als iund Achse 1 als bezeichnet j.

• Indem wir die Beschriftung in beiden Eingabearrays wiederholeni , sagen wir, einsumdass diese beiden Achsen miteinander multipliziert werden sollten. Mit anderen Worten, wir multiplizieren das Array Amit jeder Spalte des Arrays B, genau wie dies der A[:, np.newaxis] * BFall ist.

• Beachten Sie, dass jdies in unserer gewünschten Ausgabe nicht als Beschriftung angezeigt wird. Wir haben gerade verwendet i(wir möchten mit einem 1D-Array enden). Durch das Weglassen des Etiketts, doch sagen wir summieren entlang dieser Achse. Mit anderen Worten, wir summieren die Zeilen der Produkte genau wie dies der Fall ist.einsum.sum(axis=1)

Das ist im Grunde alles, was Sie wissen müssen, um zu verwenden einsum. Es hilft, ein wenig herumzuspielen; Wenn wir beide Beschriftungen in der Ausgabe belassen, erhalten 'i,ij->ij'wir ein 2D-Array von Produkten zurück (wie A[:, np.newaxis] * B). Wenn wir keine Ausgabeetiketten sagen,'i,ij-> , erhalten wir eine einzelne Zahl zurück (genau wie beim Ausführen (A[:, np.newaxis] * B).sum()).

Das Tolle daran einsum jedoch, dass nicht zuerst eine temporäre Reihe von Produkten erstellt wird. es summiert nur die Produkte, wie es geht. Dies kann zu großen Einsparungen bei der Speichernutzung führen.

Ein etwas größeres Beispiel

Um das Punktprodukt zu erklären, sind hier zwei neue Arrays:

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

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

Wir werden das Punktprodukt mit berechnen np.einsum('ij,jk->ik', A, B). Hier ist ein Bild, das die Beschriftung des Aund Bund des Ausgabearrays zeigt, die wir von der Funktion erhalten:

Sie können sehen, dass die Beschriftung jwiederholt wird. Dies bedeutet, dass wir die Zeilen von Amit den Spalten von multiplizieren B. Darüber hinaus ist das Etikett jnicht in der Ausgabe enthalten - wir fassen diese Produkte zusammen. Etiketten iundk werden für die Ausgabe aufbewahrt, sodass wir ein 2D-Array zurückerhalten.

Es könnte noch deutlicher sein , dieses Ergebnis mit dem Array zu vergleichen , wo das Etikett jwird nicht summiert. Unten links sehen Sie das 3D-Array, das sich aus dem Schreiben ergibt np.einsum('ij,jk->ijk', A, B)(dh wir haben die Beschriftung beibehalten j):

Summierachse j gibt das erwartete Punktprodukt an, das rechts gezeigt wird.

Einige Übungen

Um mehr Gefühl dafür zu bekommen einsum, kann es nützlich sein, vertraute NumPy-Array-Operationen unter Verwendung der tiefgestellten Notation zu implementieren. Alles, was Kombinationen aus Multiplikations- und Summierachsen beinhaltet, kann mit geschrieben werden einsum .

A und B seien zwei 1D-Arrays gleicher Länge. Zum Beispiel A = np.arange(10)und B = np.arange(5, 15).

• Die Summe von Akann geschrieben werden:

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

• Elementweise Multiplikation A * Bkann geschrieben werden:

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

• Das innere Produkt oder Punktprodukt np.inner(A, B)oder np.dot(A, B)kann geschrieben werden:

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

• Das äußere Produkt np.outer(A, B)kann geschrieben werden:

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

Für 2D-Arrays Cund D, vorausgesetzt, die Achsen sind kompatible Längen (beide haben dieselbe Länge oder eine von ihnen hat die Länge 1), hier einige Beispiele:

• Die Spur von C(Summe der Hauptdiagonale) np.trace(C)kann geschrieben werden:

  np.einsum('ii', C)

• Elementweise Multiplikation Cund die transponierte D, C * D.Tkann geschrieben werden:

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

• Das Multiplizieren jedes Elements Cmit dem Array D(um ein 4D-Array zu C[:, :, None, None] * Derstellen ) kann wie folgt geschrieben werden:

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

Die Idee von zu erfassen numpy.einsum()ist sehr einfach, wenn Sie es intuitiv verstehen. Beginnen wir als Beispiel mit einer einfachen Beschreibung der Matrixmultiplikation .

Um dies zu verwenden numpy.einsum(), müssen Sie lediglich die sogenannte Indexzeichenfolge als Argument übergeben, gefolgt von Ihren Eingabearrays .

Angenommen, Sie haben zwei 2D - Arrays, Aund B, und Sie wollen Matrix - Multiplikation zu tun. Also tust du es:

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

Hier entspricht die Indexzeichenfolge ij dem Array, Awährend die Indexzeichenfolge jk dem Array entspricht B. Das Wichtigste dabei ist auch, dass die Anzahl der Zeichen in jeder tiefgestellten Zeichenfolge mit den Abmessungen des Arrays übereinstimmen muss . (dh zwei Zeichen für 2D-Arrays, drei Zeichen für 3D-Arrays usw.) Wenn Sie die Zeichen zwischen tiefgestellten Zeichenfolgen ( jin unserem Fall) wiederholen , bedeutet dies, dass die einSumme entlang dieser Dimensionen erfolgen soll. Somit werden sie summenreduziert. (dh diese Dimension wird weg sein )

Die Indexzeichenfolge danach ->ist unser resultierendes Array. Wenn Sie es leer lassen, wird alles summiert und als Ergebnis ein Skalarwert zurückgegeben. Andernfalls hat das resultierende Array Dimensionen entsprechend der tiefgestellten Zeichenfolge . In unserem Beispiel wird es sein ik. Dies ist intuitiv, da wir wissen, dass für die Matrixmultiplikation die Anzahl der Spalten im Array Amit der Anzahl der Zeilen im Array übereinstimmen muss, Bwas hier geschieht (dh wir codieren dieses Wissen, indem wir das Zeichen jin der tiefgestellten Zeichenfolge wiederholen ).

Hier sind einige weitere Beispiele , die kurz und bündig die Verwendung / Leistungsfähigkeit bei der np.einsum()Implementierung einiger gängiger Tensor- oder nd-Array- Operationen veranschaulichen .

Eingänge

# 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) Matrixmultiplikation (ähnlich wie 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) Elemente entlang der Hauptdiagonale extrahieren (ähnlich wie np.diag(arr))

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

3) Hadamard-Produkt (dh elementweises Produkt zweier Arrays) (ähnlich 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) Elementweises Quadrieren (ähnlich np.square(arr)oder 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) Spur (dh Summe der Hauptdiagonalelemente) (ähnlich wie np.trace(arr))

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

6) Matrixtransponierung (ähnlich wie 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) Äußeres Produkt (von Vektoren) (ähnlich 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) Inneres Produkt (von Vektoren) (ähnlich np.inner(vec1, vec2))

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

9) Summe entlang der Achse 0 (ähnlich wie np.sum(arr, axis=0))

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

10) Summe entlang der Achse 1 (ähnlich wie np.sum(arr, axis=1))

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

11) Batch-Matrix-Multiplikation

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) Summe entlang der Achse 2 (ähnlich wie np.sum(arr, axis=2))

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

13) Summiere alle Elemente im Array (ähnlich wie np.sum(arr))

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

14) Summe über mehrere Achsen (dh Marginalisierung)
(ähnlich wie 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 Produkte (ähnlich np.sum (hadamard Produkt) s 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- und 3D-Array-Multiplikation

Eine solche Multiplikation kann sehr nützlich sein, wenn Sie ein lineares Gleichungssystem ( Ax = b ) lösen, bei dem Sie das Ergebnis überprüfen möchten.

# 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

Im Gegenteil, wenn man np.matmul()für diese Überprüfung verwenden muss, müssen wir einige reshapeOperationen ausführen, um das gleiche Ergebnis zu erzielen:

# 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 : Lesen Sie hier mehr Mathe: Einstein-Summation und definitiv hier: Tensor-Notation

Erstellen wir zwei Arrays mit unterschiedlichen, aber kompatiblen Abmessungen, um deren Zusammenspiel hervorzuheben

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

Ihre Berechnung verwendet einen 'Punkt' (Summe der Produkte) von a (2,3) mit a (3,4), um ein (4,2) -Array zu erzeugen. iist das 1. Dim von A, das letzte von C; kder letzte von B, 1. von C. jwird durch die Summe 'verbraucht'.

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

Dies ist das Gleiche wie np.dot(A,B).T- es ist die endgültige Ausgabe, die transponiert wird.

Um mehr darüber zu jerfahren, was passiert , ändern Sie die CIndizes 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]]])

Dies kann auch hergestellt werden mit:

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

Fügen Sie also eine kDimension am Ende Aund eine am Anfang ivon hinzu B, was zu einem (2,3,4) -Array führt.

0 + 4 + 16 = 20, 9 + 28 + 55 = 92Etc; Summiere jund transponiere, um das frühere Ergebnis zu erhalten:

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

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

Ich fand NumPy: Die Tricks des Handels (Teil II) lehrreich

Wir verwenden ->, um die Reihenfolge des Ausgabearrays anzugeben. Stellen Sie sich also 'ij, i-> j' mit linker Seite (LHS) und rechter Seite (RHS) vor. Jede Wiederholung von Etiketten auf der LHS berechnet das Produktelement weise und summiert sich dann. Durch Ändern der Beschriftung auf der rechten Seite (Ausgabeseite) können wir die Achse definieren, in der in Bezug auf das Eingabearray vorgegangen werden soll, dh die Summierung entlang der Achse 0, 1 usw.

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)

Beachten Sie, dass es drei Achsen gibt, i, j, k, und dass j wiederholt wird (auf der linken Seite). i,jstellen Zeilen und Spalten für dar a. j,kfür b.

Um das Produkt zu berechnen und die jAchse auszurichten , müssen wir eine Achse hinzufügen a. ( bwird entlang (?) der ersten Achse ausgestrahlt)

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

jfehlt auf der rechten Seite, also summieren wir, über jwelcher die zweite Achse des 3x3x3-Arrays liegt

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

Schließlich werden die Indizes auf der rechten Seite (alphabetisch) umgekehrt, sodass wir transponieren.

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

Beim Lesen von einsum-Gleichungen fand ich es am hilfreichsten, sie nur mental auf ihre imperativen Versionen reduzieren zu können.

Beginnen wir mit der folgenden (imposanten) Aussage:

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

Wenn wir zuerst die Interpunktion durcharbeiten, sehen wir, dass wir zwei durch Kommas getrennte 4-Buchstaben-Blobs haben - bhwi und bhwjvor dem Pfeil und einen einzelnen 3-Buchstaben-Blob bijdanach. Daher erzeugt die Gleichung ein Tensor-Ergebnis vom Rang 3 aus zwei Tensoreingaben vom Rang 4.

Lassen Sie nun jeden Buchstaben in jedem Blob den Namen einer Bereichsvariablen sein. Die Position, an der der Buchstabe im Blob erscheint, ist der Index der Achse, über die er sich in diesem Tensor erstreckt. Die imperative Summation, die jedes Element von C erzeugt, muss daher mit drei verschachtelten for-Schleifen beginnen, eine für jeden Index von 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] = ...

Sie haben also im Wesentlichen eine for Schleife für jeden Ausgabeindex von C. Wir lassen die Bereiche vorerst unbestimmt.

Als nächstes schauen wir uns die linke Seite an - gibt es dort Bereichsvariablen, die nicht auf der rechten Seite erscheinen? In unserem Fall - ja hund w. Fügen Sie eine innere Verschachtelung hinzufor für jede solche Variable Schleife hinzu:

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

In der innersten Schleife sind jetzt alle Indizes definiert, sodass wir die eigentliche Summierung schreiben können und die Übersetzung abgeschlossen ist:

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

Wenn Sie bisher dem Code folgen konnten, dann herzlichen Glückwunsch! Dies ist alles, was Sie brauchen, um einsum-Gleichungen lesen zu können. Beachten Sie insbesondere, wie die ursprüngliche Einsum-Formel der endgültigen Summationsanweisung im obigen Snippet zugeordnet ist. Die for-Schleifen und Bereichsgrenzen sind nur Flusen und diese abschließende Aussage ist alles, was Sie wirklich brauchen, um zu verstehen, was los ist.

Der Vollständigkeit halber sehen wir uns an, wie die Bereiche für jede Bereichsvariable bestimmt werden. Nun, der Bereich jeder Variablen ist einfach die Länge der Dimension (en), die sie indiziert. Wenn eine Variable mehr als eine Dimension in einem oder mehreren Tensoren indiziert, müssen natürlich die Längen jeder dieser Dimensionen gleich sein. Hier ist der obige Code mit den vollständigen Bereichen:

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