python - Entendiendo el einsum de NumPy
arrays multidimensional-array (4)
(Nota: esta respuesta se basa en una breve
publicación de blog
sobre
einsum
que escribí hace un tiempo).
¿Qué hace el
einsum
?
Imagine que tenemos dos matrices multidimensionales,
A
y
B
Ahora supongamos que queremos ...
-
multiplique
AconBde una manera particular para crear una nueva gama de productos; y luego tal vez - suma esta nueva matriz a lo largo de ejes particulares; y luego tal vez
- transponer los ejes de la nueva matriz en un orden particular.
Existe una buena posibilidad de que
einsum
nos ayude a hacer esto más rápido y con mayor eficiencia de memoria que las combinaciones de las funciones NumPy como
multiply
,
sum
y
transpose
permitirán.
¿Cómo funciona el
einsum
?
Aquí hay un ejemplo simple (pero no completamente trivial). Tome las siguientes dos matrices:
A = np.array([0, 1, 2])
B = np.array([[ 0, 1, 2, 3],
[ 4, 5, 6, 7],
[ 8, 9, 10, 11]])
Multiplicaremos
A
y
B
elementos y luego sumaremos a lo largo de las filas de la nueva matriz.
En NumPy "normal" escribiríamos:
>>> (A[:, np.newaxis] * B).sum(axis=1)
array([ 0, 22, 76])
Entonces, aquí, la operación de indexación en
A
alinea los primeros ejes de las dos matrices para que se pueda transmitir la multiplicación.
Las filas de la matriz de productos se suman para devolver la respuesta.
Ahora, si quisiéramos usar
einsum
, podríamos escribir:
>>> np.einsum(''i,ij->i'', A, B)
array([ 0, 22, 76])
La cadena de
firma
''i,ij->i''
es la clave aquí y necesita un poco de explicación.
Puedes pensarlo en dos mitades.
En el lado izquierdo (a la izquierda de
->
) hemos etiquetado las dos matrices de entrada.
A la derecha de
->
, hemos etiquetado la matriz con la que queremos terminar.
Esto es lo que sucede después:
-
Atiene un eje; lo hemos etiquetadoi. YBtiene dos ejes; hemos etiquetado el eje 0 comoiy el eje 1 comoj. -
Al repetir la etiqueta
ien ambas matrices de entrada, le estamos diciendo aleinsumque estos dos ejes deben multiplicarse juntos. En otras palabras, estamos multiplicando la matrizAcon cada columna de la matrizB, tal como lo haceA[:, np.newaxis] * B -
Observe que
jno aparece como una etiqueta en nuestra salida deseada; acabamos de usari(queremos terminar con una matriz 1D). Al omitir la etiqueta, le estamos diciendo aleinsumqueeinsuma lo largo de este eje. En otras palabras, estamos sumando las filas de los productos, tal como lo hace.sum(axis=1).
Eso es básicamente todo lo que necesita saber para usar
einsum
.
Ayuda a jugar un poco;
si dejamos ambas etiquetas en la salida,
''i,ij->ij''
, recuperamos una matriz 2D de productos (igual que
A[:, np.newaxis] * B
).
Si decimos que no hay etiquetas de salida,
''i,ij->
, recuperamos un solo número (lo mismo que hacer
(A[:, np.newaxis] * B).sum()
).
einsum
embargo, lo mejor de
einsum
es que no crea primero una gama temporal de productos;
simplemente resume los productos a medida que avanza.
Esto puede conducir a grandes ahorros en el uso de la memoria.
Un ejemplo un poco más grande
Para explicar el producto punto, aquí hay dos nuevas matrices:
A = array([[1, 1, 1],
[2, 2, 2],
[5, 5, 5]])
B = array([[0, 1, 0],
[1, 1, 0],
[1, 1, 1]])
Calcularemos el producto punto usando
np.einsum(''ij,jk->ik'', A, B)
.
Aquí hay una imagen que muestra el etiquetado de
A
y
B
y la matriz de salida que obtenemos de la función:
Puede ver que la etiqueta
j
se repite; esto significa que estamos multiplicando las filas de
A
con las columnas de
B
Además, la etiqueta
j
no está incluida en la salida; estamos sumando estos productos.
Las etiquetas
i
y
k
se guardan para la salida, por lo que recuperamos una matriz 2D.
Puede ser aún más claro comparar este resultado con la matriz donde
no se
suma la etiqueta
j
.
A continuación, a la izquierda, puede ver la matriz 3D que resulta de escribir
np.einsum(''ij,jk->ijk'', A, B)
(es decir, hemos mantenido la etiqueta
j
):
Sumar el eje
j
da el producto puntual esperado, que se muestra a la derecha.
Algunos ejercicios
Para tener una mejor idea de
einsum
, puede ser útil implementar operaciones familiares de matriz NumPy utilizando la notación de subíndice.
Cualquier cosa que implique combinaciones de ejes de multiplicación y suma se puede escribir usando
einsum
.
Deje que A y B sean dos matrices 1D con la misma longitud.
Por ejemplo,
A = np.arange(10)
y
B = np.arange(5, 15)
.
-
La suma de
Ase puede escribir:np.einsum(''i->'', A) -
La multiplicación por elementos,
A * B, se puede escribir:np.einsum(''i,i->i'', A, B) -
El producto interno o producto
np.inner(A, B),np.inner(A, B)onp.dot(A, B), se puede escribir:np.einsum(''i,i->'', A, B) # or just use ''i,i'' -
El producto externo,
np.outer(A, B), se puede escribir:np.einsum(''i,j->ij'', A, B)
Para las matrices 2D,
C
y
D
, siempre que los ejes sean longitudes compatibles (tanto la misma longitud como una de ellas tiene longitud 1), aquí hay algunos ejemplos:
-
La traza de
C(suma de la diagonal principal),np.trace(C), se puede escribir:np.einsum(''ii'', C) -
La multiplicación por elementos de
Cy la transposición deD,C * DT, se pueden escribir:np.einsum(''ij,ji->ij'', C, D) -
Multiplicando cada elemento de
Cpor la matrizD(para hacer una matriz 4D),C[:, :, None, None] * D, se puede escribir:np.einsum(''ij,kl->ijkl'', C, D)
Estoy luchando por entender exactamente cómo funciona el
einsum
.
He visto la documentación y algunos ejemplos, pero no parece que se quede.
Aquí hay un ejemplo que analizamos en clase:
C = np.einsum("ij,jk->ki", A, B)
para dos matrices
A
y
B
Creo que esto tomaría
A^T * B
, pero no estoy seguro (está tomando la transposición de uno de ellos, ¿verdad?).
¿Alguien puede
einsum
exactamente qué está sucediendo aquí (y en general cuando se usa
einsum
)?
Agarrar la idea de
numpy.einsum()
es muy fácil si lo entiendes intuitivamente.
Como ejemplo, comencemos con una descripción simple que involucra
la multiplicación de matrices
.
Para usar
numpy.einsum()
, todo lo que tiene que hacer es pasar la llamada
cadena de subíndices
como argumento, seguida de sus
matrices de entrada
.
Digamos que tiene dos matrices 2D,
A
y
B
, y desea hacer una multiplicación de matrices.
Tu también:
np.einsum("ij, jk -> ik", A, B)
Aquí la
cadena de subíndice
ij
corresponde a la matriz
A
mientras que la
cadena de subíndice
jk
corresponde a la matriz
B
Además, lo más importante a tener en cuenta aquí es que el
número de caracteres
en cada
cadena de subíndice
debe
coincidir con las dimensiones de la matriz.
(es decir, dos caracteres para matrices 2D, tres caracteres para matrices 3D, etc.) Y si repite los caracteres entre
cadenas de subíndice
(
j
en nuestro caso), eso significa que quiere que la
suma de
ein
suceda a lo largo de esas dimensiones.
Por lo tanto, se reducirán la suma.
(es decir, esa dimensión se habrá
ido
)
La
cadena de subíndice
después de esto
->
, será nuestra matriz resultante.
Si lo deja vacío, todo se sumará y se devolverá un valor escalar como resultado.
De lo contrario, la matriz resultante tendrá dimensiones de acuerdo con la
cadena de subíndice
.
En nuestro ejemplo, será
ik
.
Esto es intuitivo porque sabemos que para la multiplicación de matrices, el número de columnas en la matriz
A
tiene que coincidir con la cantidad de filas en la matriz
B
que es lo que está sucediendo aquí (es decir, codificamos este conocimiento repitiendo el char
j
en la
cadena del subíndice
)
Aquí hay algunos ejemplos más que ilustran el uso / poder de numpy.einsum() en la implementación de algunas operaciones comunes de tensor o nd-array , de manera sucinta.
Entradas
# 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) Multiplicación matricial
(similar 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) Extraer elementos a lo largo de la diagonal principal
(similar a
np.diag(arr)
)
In [202]: np.einsum("ii -> i", A)
Out[202]: array([11, 22, 33, 44])
3) Producto Hadamard (es decir, producto
arr1 * arr2
elementos de dos matrices)
(similar 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)
np.square(arr)
elementos
(similar 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) Trace (es decir, suma de elementos de diagonal principal)
(similar a
np.trace(arr)
)
In [217]: np.einsum("ii -> ", A)
Out[217]: 110
6) Transposición de matriz
(similar 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) Producto externo (de vectores)
(similar 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) Producto interno (de vectores)
(similar a
np.inner(vec1, vec2)
)
In [256]: np.einsum("i, i -> ", vec, vec)
Out[256]: 14
9) Suma a lo largo del eje 0
(similar a
np.sum(arr, axis=0)
)
In [260]: np.einsum("ij -> j", B)
Out[260]: array([10, 10, 10, 10])
10) Suma a lo largo del eje 1
(similar a
np.sum(arr, axis=1)
)
In [261]: np.einsum("ij -> i", B)
Out[261]: array([ 4, 8, 12, 16])
11) Multiplicación de matriz de lotes
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) Suma a lo largo del eje 2
(similar 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) Suma todos los elementos en la matriz
(similar a
np.sum(arr)
)
In [335]: np.einsum("ijk -> ", BM)
Out[335]: 480
14) Suma sobre múltiples ejes (es decir, marginación)
(similar 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) Productos de doble punto (similar a 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) Multiplicación de matrices 2D y 3D
Tal multiplicación podría ser muy útil cuando se resuelve un sistema lineal de ecuaciones ( Ax = b ) donde desea verificar el resultado.
# 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
Por el contrario, si uno tiene que usar
np.matmul()
para esta verificación, tenemos que hacer un par de operaciones de
reshape
para lograr el mismo resultado como:
# 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
Bono : Lea más matemáticas aquí: Einstein-Summation y definitivamente aquí: Tensor-Notation
Encontré NumPy: Los trucos del oficio (Parte II) instructivos
Usamos -> para indicar el orden de la matriz de salida. Así que piense que ''ij, i-> j'' tiene el lado izquierdo (LHS) y el lado derecho (RHS). Cualquier repetición de etiquetas en el LHS calcula el elemento del producto sabiamente y luego suma. Al cambiar la etiqueta en el lado RHS (salida), podemos definir el eje en el que queremos proceder con respecto a la matriz de entrada, es decir, la suma a lo largo del eje 0, 1 y así sucesivamente.
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)
Observe que hay tres ejes, i, j, k, y que j se repite (en el lado izquierdo).
i,j
representan filas y columnas para
a
.
j,k
para
b
.
Para calcular el producto y alinear el eje
j
, necesitamos agregar un eje a
a
.
(
b
se transmitirá a lo largo (?) del primer eje)
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á ausente del lado derecho, por lo que sumamos sobre
j
que es el segundo eje de la matriz 3x3x3
>>> c = c.sum(1)
>>> c
array([[ 9, 12, 15],
[18, 24, 30],
[27, 36, 45]])
Finalmente, los índices se invierten (alfabéticamente) en el lado derecho, por lo que transponemos.
>>> 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]])
>>>
Hagamos 2 matrices, con dimensiones diferentes pero compatibles para resaltar su interacción
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]])
Su cálculo, toma un ''punto'' (suma de productos) de un (2,3) con un (3,4) para producir una matriz (4,2).
i
es el primer dim de
A
, el último de
C
;
k
el último de
B
, primero de
C
j
es ''consumido'' por la sumatoria.
In [45]: C=np.einsum(''ij,jk->ki'',A,B)
Out[45]:
array([[20, 56],
[23, 68],
[26, 80],
[29, 92]])
Esto es lo mismo que
np.dot(A,B).T
: es la salida final que se transpone.
Para ver más de lo que le sucede a
j
, cambie los subíndices
C
a
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]]])
Esto también se puede producir con:
A[:,:,None]*B[None,:,:]
Es decir, agregue una dimensión
k
al final de
A
y una
i
al frente de
B
, lo que da como resultado una matriz (2,3,4).
0 + 4 + 16 = 20
,
9 + 28 + 55 = 92
, etc.
Suma en
j
y transpone para obtener el resultado anterior:
np.sum(A[:,:,None] * B[None,:,:], axis=1).T
# C[k,i] = sum(j) A[i,j (,k) ] * B[(i,) j,k]