web-dev-qa-db-fra.com

Comprendre l'einsum de NumPy

J'ai du mal à comprendre exactement comment einsum fonctionne. J'ai regardé la documentation et quelques exemples, mais cela ne semble pas coller.

Voici un exemple que nous avons parcouru en classe:

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

pour deux tableaux A et B

Je pense que cela prendrait A^T * B, mais je ne suis pas sûr (cela prend la transposition de l’un d’eux, non?). Quelqu'un peut-il m'expliquer exactement ce qui se passe ici (et en général lors de l'utilisation de einsum)?

147
Lance Strait

(Remarque: cette réponse est basée sur un court message message de blog à propos de einsum que j'ai écrit il y a quelque temps.)

Que fait einsum?

Imaginons que nous ayons deux tableaux multidimensionnels, A et B. Supposons maintenant que nous voulons ...

  • multiplierA avec B d'une manière particulière pour créer un nouveau tableau de produits; et puis peut-être
  • sum ce nouveau tableau suivant des axes particuliers; et puis peut-être
  • transpose les axes du nouveau tableau dans un ordre particulier.

Il y a de bonnes chances que einsum nous aide à faire cela plus rapidement et plus efficacement en mémoire que des combinaisons des fonctions NumPy comme multiply, sum et transpose Autoriser.

Comment fonctionne einsum?

Voici un exemple simple (mais pas complètement trivial). Prenez les deux tableaux suivants:

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

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

Nous multiplierons A et B par élément, puis additionnerons le long des lignes du nouveau tableau. En NumPy "normal" nous écririons:

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

Donc, ici, l’opération d’indexation sur A aligne les premiers axes des deux tableaux afin que la multiplication puisse être diffusée. Les lignes du tableau de produits sont ensuite additionnées pour renvoyer la réponse.

Maintenant, si nous voulions utiliser einsum à la place, nous pourrions écrire:

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

Le signature la chaîne 'i,ij->i' Est la clé ici et nécessite quelques explications. Vous pouvez y penser en deux moitiés. Sur le côté gauche (à gauche du ->), Nous avons étiqueté les deux tableaux d'entrée. À la droite de ->, Nous avons étiqueté le tableau que nous voulons finir.

Voici ce qui se passe ensuite:

  • A a un axe; nous l'avons étiqueté i. Et B a deux axes; nous avons étiqueté l’axe 0 comme i et l’axe 1 comme j.

  • Par en répétant ​​le libellé i dans les deux tableaux d'entrée, nous disons einsum que ces deux axes doivent être multipliés ensemble. En d'autres termes, nous multiplions le tableau A avec chaque colonne du tableau B, comme le fait A[:, np.newaxis] * B.

  • Notez que j n'apparaît pas comme une étiquette dans notre sortie désirée; nous venons d'utiliser i (nous voulons finir avec un tableau 1D). En omettant ​​l'étiquette, nous disons einsum à somme le long de cet axe. En d'autres termes, nous additionnons les lignes des produits, comme le fait .sum(axis=1).

C'est tout ce que vous devez savoir pour utiliser einsum. Ça aide de jouer un peu; si nous laissons les deux étiquettes dans la sortie, 'i,ij->ij', nous récupérons un tableau 2D de produits (identique à A[:, np.newaxis] * B). Si nous disons aucune étiquette de sortie, 'i,ij->, Nous récupérons un numéro unique (comme si (A[:, np.newaxis] * B).sum()).

Le point positif de einsum est toutefois qu’il ne crée pas d’abord un tableau temporaire de produits; il résume simplement les produits au fur et à mesure. Cela peut entraîner d'importantes économies d'utilisation de la mémoire.

Un exemple un peu plus grand

Pour expliquer le produit scalaire, voici deux nouveaux tableaux:

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

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

Nous allons calculer le produit scalaire en utilisant np.einsum('ij,jk->ik', A, B). Voici une image montrant l'étiquetage de A et B et du tableau de sortie obtenu à partir de la fonction:

enter image description here

Vous pouvez voir que l'étiquette j est répétée - cela signifie que nous multiplions les lignes de A par les colonnes de B. De plus, l'étiquette j n'est pas incluse dans la sortie - nous additionnons ces produits. Les étiquettes i et k sont conservées pour la sortie, nous récupérons donc un tableau 2D.

Il peut être encore plus clair de comparer ce résultat avec le tableau où l’étiquette j est pas sommée. Ci-dessous, à gauche, vous pouvez voir le tableau 3D résultant de l'écriture de np.einsum('ij,jk->ijk', A, B) (c'est-à-dire que nous avons conservé le libellé j):

enter image description here

L'axe des sommations j donne le produit scalaire attendu, illustré à droite.

Quelques exercices

Pour obtenir plus de sensations pour einsum, il peut être utile d'implémenter des opérations de tableau NumPy familières à l'aide de la notation en indice. Tout ce qui implique des combinaisons d'axes de multiplication et de sommation peut être écrit en utilisant einsum.

Soit A et B deux matrices 1D de même longueur. Par exemple, A = np.arange(10) et B = np.arange(5, 15).

  • La somme de A peut s'écrire:

    np.einsum('i->', A)
    
  • La multiplication par élément, A * B, Peut s'écrire:

    np.einsum('i,i->i', A, B)
    
  • Le produit interne ou produit scalaire, np.inner(A, B) ou np.dot(A, B), peut être écrit:

    np.einsum('i,i->', A, B) # or just use 'i,i'
    
  • Le produit externe, np.outer(A, B), peut être écrit:

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

Pour les tableaux 2D, C et D, à condition que les axes soient des longueurs compatibles (la même longueur ou l'un d'entre eux a la longueur 1), en voici quelques exemples:

  • La trace de C (somme de la diagonale principale), np.trace(C), peut s’écrire:

    np.einsum('ii', C)
    
  • La multiplication élémentaire de C et la transposition de D, C * D.T, Peuvent s'écrire:

    np.einsum('ij,ji->ij', C, D)
    
  • En multipliant chaque élément de C par le tableau D (pour créer un tableau 4D), C[:, :, None, None] * D, Vous écrivez:

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

Saisir l'idée de numpy.einsum() est très facile si vous le comprenez de manière intuitive. A titre d'exemple, commençons par une description simple impliquant multiplication de matrice.


Pour utiliser numpy.einsum() , tout ce que vous avez à faire est de passer le soi-disant subscripts string en argument, suivi de votre tableaux d'entrée.

Disons que vous avez deux tableaux 2D, A et B, et vous voulez faire la multiplication de matrice. Alors tu fais:

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

Ici le chaîne d'indiceij correspond à un tableau A avec le chaîne d'indicejk correspond au tableau B. En outre, la chose la plus importante à noter ici est que le nombre de caractères dans chaque chaîne d'indice doit correspondent aux dimensions du tableau. (c.-à-d. deux caractères pour les tableaux 2D, trois caractères pour les tableaux 3D, etc.). Et si vous répétez les caractères entre chaînes d'indice (j dans notre cas), cela signifie que vous voulez que ein somme se produise selon ces dimensions. Ainsi, ils seront réduits en somme. (c’est-à-dire que cette dimension sera disparu)

Le chaîne d'indice après cela -> sera notre tableau résultant. Si vous le laissez vide, tout sera résumé et une valeur scalaire sera renvoyée comme résultat. Sinon, le tableau résultant aura des dimensions conformes à la chaîne d'indice. Dans notre exemple, ce sera ik. Ceci est intuitif car nous savons que pour la multiplication de matrice, le nombre de colonnes dans un tableau A doit correspondre au nombre de lignes dans un tableau B qui correspond à ce qui se passe ici (c’est-à-dire que nous encodons cette connaissance en répétant le caractère j dans la chaîne indice)


Voici quelques exemples supplémentaires illustrant l'utilisation/la puissance de np.einsum() dans la mise en œuvre de certains tensor ou nd-array opérations, succinctement.

Entrées

# 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) Multiplication matricielle (similaire à 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) Extrait des éléments le long de la diagonale principale (semblable à np.diag(arr))

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

3) Produit Hadamard (produit élémentaire de deux tableaux) (similaire à 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) Carré d'élément (similaire à np.square(arr) ou 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) Trace (c'est-à-dire la somme des éléments de la diagonale principale) (similaire à np.trace(arr))

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

6) Transposition de matrice (similaire à 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) Produit extérieur (de vecteurs) (similaire à 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) Produit intérieur (de vecteurs) (similaire à np.inner(vec1, vec2))

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

9) Somme le long de l'axe 0 (similaire à np.sum(arr, axis=0))

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

10) Somme le long de l’axe 1 (similaire à np.sum(arr, axis=1))

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

11) Multiplication de matrice par lots

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) Somme le long de l'axe 2 (similaire à np.sum(arr, axis=2))

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

13) Somme tous les éléments d'un tableau (similaire à np.sum(arr))

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

14) Somme sur plusieurs axes (marginalisation)
(similaire à 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) produits à double point (similaire à 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

16) Multiplication de tableaux 2D et 3D

Une telle multiplication peut s'avérer très utile lors de la résolution d'un système d'équations linéaire ( Ax = b ) où vous souhaitez vérifier le résultat.

# 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

Au contraire, si on doit utiliser np.matmul() pour cette vérification, il faut faire un couple de reshape opérations pour obtenir le même résultat comme:

# 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 : Lisez plus de mathématiques ici: Einstein-Summation et définitivement ici: Tensor-Notation

33
kmario23

Permet de faire 2 tableaux, avec des dimensions différentes, mais compatibles pour mettre en évidence leur interaction

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

Votre calcul prend un "point" (somme de produits) d'un (2,3) avec un (3,4) pour produire un tableau (4,2). i est la 1ère dim de A, la dernière de C; k le dernier de B, le premier de C. j est "consommé" par la somme.

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

C'est la même chose que np.dot(A,B).T - c'est la sortie finale qui est transposée.

Pour en savoir plus sur ce qui arrive à j, remplacez les indices C par 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]]])

Ceci peut également être produit avec:

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

C'est-à-dire que vous ajoutez une dimension k à la fin de A, et un i au début de B, ce qui donne un (2,3, 4) tableau.

0 + 4 + 16 = 20, 9 + 28 + 55 = 92, Etc; Faites la somme sur j et transposez-la pour obtenir le résultat précédent:

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

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

J'ai trouvé NumPy: Les ficelles du métier (partie II) instructif

Nous utilisons -> pour indiquer l’ordre du tableau de sortie. Donc, pensez à 'ij, i-> j' comme ayant le côté gauche (LHS) et le côté droit (RHS). Toute répétition d’étiquettes sur le LHS calcule l’élément de produit, puis effectue la somme. En modifiant l’étiquette du côté RHS (sortie), nous pouvons définir l’axe dans lequel nous voulons procéder par rapport au tableau d’entrée, c’est-à-dire la somme le long des axes 0, 1, etc.

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)

Remarquez qu'il y a trois axes, i, j, k et que j est répété (du côté gauche). i,j représente les lignes et les colonnes de a. j,k pour b.

Afin de calculer le produit et d’aligner l’axe j, nous devons ajouter un axe à a. (b sera diffusé le long du premier axe)

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 est absent du côté droit de sorte que nous somme sur j qui est le deuxième axe du tableau 3x3x3

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

Enfin, les index sont inversés (par ordre alphabétique) à droite, de sorte que nous les transposons.

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