es.davy.ai

Preguntas y respuestas de programación confiables

¿Tienes una pregunta?

Si tienes alguna pregunta, puedes hacerla a continuación o ingresar lo que estás buscando.

Cómo indexar un tensor con un tensor de alta dimensión.

28 May, 2023 Programación

Tengo dos tensores, x para valores y y para indexación. x.shape e y.shape son iguales excepto por la última dimensión.
Por ejemplo:

x=torch.tensor([[1, 6, 7, 5, 6],
        [8, 6, 7, 8, 4],
        [2, 8, 3, 5, 6]])

<h1>x.shape:(3,5)</h1>

y=torch.tensor([[1, 2],
        [2, 3],
        [2, 2]])

<h1>y.shape:(3,2)</h1>

<code>
¿Hay una manera simple de hacer slice de `x[y]` para que el resultado sea:
</code>
torch.tensor([6,7],[7,8],[3,3])

¿Qué pasa si x e y son tensores de dimensiones superiores?:

“`
x.shape=(a,b,c,d)
y.shape=(a,b,c,e)

<h1>a, b, c, d, e son enteros positivos</h1>

“`

Tags:  , ,

Answer

  1. Avatar for davy.ai
    davy.ai

    Sí, hay una manera simple de cortar x[y].

    Podemos usar el indexado avanzado en PyTorch para lograr esto. Podemos crear una tupla de índices con la ayuda de la función torch.meshgrid y pasar esta tupla como argumento de indexación para el tensor x.

    El código para cortar x[y] para el ejemplo dado es:

    indices = torch.meshgrid(torch.arange(y.shape[0]), torch.arange(y.shape[1]))
result = x[y[indices]]

    Esto nos dará el tensor de salida deseado:

    tensor([[6, 7],
        [7, 8],
        [3, 3]])

    Si x e y son tensores de dimensiones superiores con formas (a, b, c, d) y (a, b, c, e) respectivamente, podemos aplicar la misma lógica primero creando una tupla de índices para la última dimensión:

    last_indices = torch.meshgrid(*[torch.arange(size) for size in y.shape[-2:]])

# Para tensores 3D, esto convierte last_indices en una tupla:
# (tensor([[[0, 0], [1, 1]], [[0, 0], [1, 1]], [[0, 0], [1, 1]]]),
#  tensor([[[0, 1], [0, 1]], [[0, 1], [0, 1]], [[0, 1], [0, 1]]]))

all_indices = last_indices + (y,) # concatenar los índices de la última dimensión y `y`

result = x[all_indices]

    Esto nos dará un tensor de forma (a, b, c, e) como salida.

Comments are closed.