En algunas situaciones, necesitará realizar una indexación/selección avanzada con Pytorch, por ejemplo, responder la pregunta: “¿cómo puedo seleccionar elementos del Tensor A siguiendo los índices especificados en el Tensor B?”
En esta publicación presentaremos los tres métodos más comunes para tales tareas, a saber antorcha.index_select, antorcha.reunir y antorcha.tomar. Los explicaremos todos en detalle y los contrastaremos entre sí.
Es cierto que una motivación para esta publicación fue que olvidé cómo y cuándo usar qué función, terminé buscando en Google, navegando Desbordamiento de pila y la documentación oficial, en mi opinión, relativamente breve y no demasiado útil. Por lo tanto, como se mencionó, aquí profundizamos en estas funciones: motivamos cuándo usarlas, damos ejemplos en 2 y 3D y mostramos la selección resultante gráficamente.
Espero que esta publicación aporte claridad sobre dichas funciones y elimine la necesidad de seguir explorando. ¡Gracias por leer!
Y ahora, sin más, profundicemos en las funciones una por una. Para todos, primero comenzamos con un ejemplo 2D y visualizamos la selección resultante, y luego pasamos a un ejemplo algo más complejo en 3D. Además, volvemos a implementar la operación ejecutada en Python simple; puede consultar el pseudocódigo como otra fuente de información sobre lo que hacen estas funciones. Al final, resumimos las funciones y sus diferencias en una tabla.
antorcha.index_select selecciona elementos a lo largo de una dimensión, mientras mantiene los demás sin cambios. Es decir: mantenga todos los elementos de todas las demás dimensiones, pero seleccione elementos en las dimensiones de destino siguiendo el tensor de índice. Demostremos esto con un ejemplo 2D, en el que seleccionamos a lo largo de la dimensión 1:
num_picks = 2values = torch.rand((len_dim_0, len_dim_1))
indices = torch.randint(0, len_dim_1, size=(num_picks,))
# [len_dim_0, num_picks]
picked = torch.index_select(values, 1, indices)
El tensor resultante tiene forma [len_dim_0, num_picks]: para cada elemento a lo largo de la dimensión 0, hemos elegido el mismo elemento de la dimensión 1. Visualicemos esto: