Función PyTorch torch.take

Pytorch torch 参考手册Manual de referencia de Pytorch torch

torch.takeEs una función en PyTorch que se utiliza para obtener los elementos en las posiciones de índice dadas. Trata el tensor de entrada como un arreglo unidimensional y devuelve los elementos en las posiciones de índice especificadas.

Definición de la función

torch.take(input, index)

Parámetros:

  • input(Tensor): Tensor de entrada.
  • index(Tensor): Tensor de índices enteros, especifica las posiciones de los elementos a obtener.

Valor de retorno:

  • torch.Tensor: Devuelve un nuevo tensor compuesto por los elementos en las posiciones de índice especificadas.

Ejemplo de uso

Ejemplo

import torch

# Crear un tensor 2D
x = torch.tensor([[1, 2, 3],
                  [4, 5, 6],
                  [7, 8, 9]])

print("Tensor original:")
print(x)

# Tratar el tensor 2D como un arreglo unidimensional, índices 0-8
# Obtener los elementos en las posiciones de índice 0, 4, 8
index = torch.tensor([0, 4, 8])
y = torch.take(x, index)

print("n índice:", index)
print("Elementos extraídos:", y)

El resultado de salida es:

原始张量:
tensor([[1, 2, 3],
        [4, 5, 6],
        [7, 8, 9]])

索引: tensor([0, 4, 8])
取出的元素: tensor([1, 5, 9])

Ejemplo

import torch

# Tensor 3D
x = torch.arange(24).reshape(2, 3, 4)

print("Tensor 3D original:")
print(x)
print("Forma:", x.shape)

# Obtener elementos de múltiples índices
index = torch.tensor([0, 1, 2, 10, 20, 23])
y = torch.take(x, index)

print("n índice:", index)
print("Elementos extraídos:", y)

El resultado de salida es:

原始3D张量:
tensor([[[ 0,  1,  2,  3],
         [ 4,  5,  6,  7],
         [ 8, 9, 10, 11]],

        [[12, 13, 14, 15],
         [16, 17, 18, 19],
         [20, 21, 22, 23]]])
形状: torch.Size([2, 3, 4])

索引: tensor([ 0,  1,  2, 10, 20, 23])
取出的元素: tensor([ 0,  1,  2, 10, 20, 23])

Ejemplo

import torch

# Usar índices negativos
x = torch.tensor([[1, 2, 3],
                  [4, 5, 6]])

# Los índices negativos se cuentan desde el final
index = torch.tensor([0, -1])  # El primer y el último elemento
y = torch.take(x, index)

print("Original:", x)
print("Índices [0, -1]:", y)

El resultado de salida es:

原始: tensor([[1, 2, 3],
        [4, 5, 6]])
索引 [0, -1]: tensor([1, 6])

Ejemplo

import torch

# Seleccionar elementos aleatoriamente
x = torch.randn(10, 10)

# Generar aleatoriamente 10 índices
index = torch.randint(0, 100, (10,))
print("Índices aleatorios:", index)

# Obtener los elementos en las posiciones correspondientes
selected = torch.take(x, index)

print("Forma del tensor original:", x.shape)
print("Forma de los elementos seleccionados:", selected.shape)
print("Elementos seleccionados:", selected)

El resultado de salida es:

随机索引: tensor([12, 45, 67, 82, 35, 59, 92,  7, 28, 73])
原始张量形状: torch.Size([10, 10])
选中的元素形状: torch.Size([10])
选中的元素: tensor([ 0.2345, -0.1234,  0.5678,  1.2345, -0.6789,  0.8901, -0.3456,  0.1234, -0.5678,  0.7890])

Nota:torch.takeEl tensor de entrada se indexa como un tensor unidimensional aplanado. Los índices deben estar dentro del rango válido (0 a numel-1), también se pueden usar índices negativos (-1 representa el último elemento).

Pytorch torch 参考手册Manual de referencia de Pytorch torch

Otras extensiones