Einops: Haciendo La Vida Mejor (Generalmente)

A veces cuando usas PyTorch, encontrás a escenarios en que necesitás manipular tensores de dimensiones multiples a otras formas o combinar unas dimensiones a una. Ya existe una gran cantidad de funciones específicamente para hacer esto. Por ejemplo, se ve view, permutate, stack, tile y concat muy frecuentemente. Sin embargo, un gran problema que tienen estas funciones es que a veces es muy difícil entender lo que esta pasando a los tensores y como será la forma nueva después de usarlas.

Por eso existe Einops. Este modulo fue introducido a me por un amigo también trabajando en visión por computadora, y es muy claro como puede ayudarte escribir código para manipular a tensores en una forma que es más fácil entender, y con suerte un poco más claro a la gente que tiene que leer el código después.

Para una guía un poco más intensiva, podes ver su sitio de web. En este articulo, yo voy a cobrar unos casos claros que veo para visión por computadora, y después hablar de unos casos en que no recomiendo usar Einops en lugar de las funciones ofrecido por PyTorch.

Al primero, hacemos un data loader simple para ver unos ejemplos reales.

import torch
from torchvision import datasets
from torchvision.transforms import ToTensor
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt

from einops import rearrange, reduce, repeat, pack, unpack

training_data = datasets.FashionMNIST(
    root="data",
    train=True,
    download=True,
    transform=ToTensor()
)

train_dataloader = DataLoader(training_data, batch_size=16, shuffle=True)

images, labels = next(iter(train_dataloader))

Unos Casos de Uso Potenciales

Organizando Imágenes en Un Stack

Después de experimentar con Einops un poco, hay unos situaciones que creo que serían muy útiles a mis proyectos directamente. El primero situación es organizar imágenes a una pila (stack). Hay una variedad de razones para hacer una pila, pero un caso muy simple es visualizar un batch de imágenes.  

Para mí, una persona que usa PyTorch regularmente, mi primer instinto es usar torch.stack o torch.concat. Sin embargo, si lo hacés, te vas a acordar que los dos necesitan una lista o un tule de tensores. No funcionan con un tensor solo. En lugar de eso, necesitás usar torch.view. El código sería como esto:

stacked_image = images.view(images.shape[0] * images.shape[2], images.shape[3])

figure = plt.figure(figsize=(10, 10))
plt.imshow(stacked_image.squeeze(), cmap="gray")

Como ves, es un poco difícil entender. Tenés que multiplicar valores diferentes del shape del tensor para calcular la forma nueva que querés. Lo podés simplificar un poco y definir variables nuevas para cada parte del shape, pero eso también es un poco incómodo. No es simple recordar como view funciona directamente. Si usas einops, podes asignar un nombre intuitivo a las dimensiones y también hacerlo todo adentro de una función. Es claro en el código que estas "rearranging" los tensores.

stacked_image = rearrange(images, 'b c h w -> c (b h) w')

figure = plt.figure(figsize=(10, 10))
plt.imshow(stacked_image.squeeze(), cmap="gray")

Manipulando a Bounding Boxes

Otra situación que veo mucho es combinando y separando bounding boxes y predicciones. Generalmente, si fuera yo, usaría torch.concat o torch.join para combinarlos. Pero, torch.concat require que se especifica la dimension en que queres combinar y a veces puede ser un poco difícil imaginar como los datos estarán conectado.

torch.concat([bounding_boxes, predictions.unsqueeze(dim=1)], dim=1)

En este caso, se puede usar el método "pack" de einops en lugar de eso. Usando la notación inspirado por einum , podes ver como se mantiene la altura cuando los anchos de las predicciones y bounding boxes están combinados. Para mí, es un poco difícil decir si es más fácil leer que torch.concat, pero una cosa muy útil de de "pack" es que devuelve un variable que se dice "Packed Shapes". "Packed Shapes" es fundamentalmente la información de como los datos fueron combinado.

bbox_and_pred, ps = pack([bounding_boxes, predictions], 'h *')

Esto puede ser muy útil si tenés que combinar y separar los mismos datos muchas veces. Para separar mis bounding boxes y predicciones, tendría que pensar en los indices de los tensores, pero con "unpack", puedo pasar el variable Packed Shapes y separar los datos a las formas originales muy simplemente.  

bounding_boxes, predictions = unpack(bbox_and_pred, ps, 'h *')

La unica cosa mala que yo veo de todo esto es que tenés que manejar Packed Shapes después. Por ejemplo, si queres devolver los datos combinados de una función y intentas separar los datos en otro parte de la programa, vas a necesitar devolver Packed Shapes también. Si no, no recibís los beneficios de usar la función "pack". Aunque al fin, no creo sea un problema muy grande.

Situaciones En Que No Usaría Einops

Además de todos los ejemplos buenos en la documentación de einopstambién existen unos ejemplos como MaxPool y aumentación de imágenes. Siendo honesto, aunque creo que funciones como einops "rearrange" son más claro que torch.concat o torch.view, no creo que usar una función de einops en lugar de MaxPool es mas fácil entender.

Recomiendo a cualquiera persona usando einops que piensa un rato sobre cual función es verdaderamente más fácil leer o entender. En muchos casos, "rearrange" y "pack" son mucho más intuitivos que los nombres y usos de las funciónes PyTorch. Sin embargo, como en todos idiomas de programming, tratando de escribir código demasiado ingenioso muchas veces resuelta en código imposible leer.

Source Code

Como siempre, el source code de mi post esta aquí. Esta vez solamente hice un jupyter notebook así que estoy ilustrando unas cosas básicas de las funciones de einops. einops functions.

Conclusión

Einops es un modulo muy bueno para la manipulación de tensores que puede hacer tu código más fácil leer y ahorrar tiempo cuando estás tratando de entender como manipular a tensores. Pero, no dejas de recordar que ser racional cuando lo estás usando, y siempre usa la forma más simple en lugar de la forma más ingeniosa.