Función torch.argwhere de PyTorch

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

torch.argwhereEs una función en PyTorch que se utiliza para devolver los índices de los elementos que cumplen una condición. Devuelve los índices de los elementos cuyo valor es True (distinto de cero) en el tensor de entrada.

Definición de la función

torch.argwhere(input)

Parámetros:

  • input(Tensor): El tensor de entrada.

Valor de retorno:

  • torch.Tensor: Devuelve un tensor bidimensional, donde cada fila es el índice de un elemento que cumple la condición.

Ejemplo de uso

Ejemplo

import torch

# Crear un tensor
x = torch.tensor([[1, 0, 2],
                  [0, 3, 0],
                  [4, 5, 0]])

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

# Devolver los índices de los elementos distintos de cero
indices = torch.argwhere(x)

print("nÍndices de los elementos distintos de cero:")
print(indices)

El resultado de salida es:

原始张量:
tensor([[1, 0, 2],
        [0, 3, 0],
        [4, 5, 0]])

非零元素的索引:
tensor([[0, 0],
        [0, 2],
        [1, 1],
        [2, 0],
        [2, 1]])

Ejemplo

import torch

# Condición booleana
x = torch.tensor([[True, False, True],
                  [False, True, False],
                  [True, True, False]])

print("Tensor booleano:")
print(x)

indices = torch.argwhere(x)
print("nÍndices de valores True:")
print(indices)

El resultado de salida es:

布尔张量:
tensor([[True, False, True],
        [False, True, False],
        [True, True, False]])

True 值的索引:
tensor([[0, 0],
        [0, 2],
        [1, 1],
        [2, 0],
        [2, 1]])

Ejemplo

import torch

# Encontrar elementos mayores que un valor
x = torch.randn(3, 4)
threshold = 0

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

# Encontrar los índices de los elementos mayores que threshold
indices = torch.argwhere(x > threshold)

print(f"nÍndices de elementos mayores que {threshold}:")
print(indices)

# También se puede usar la función nonzero, con el mismo efecto
indices2 = torch.nonzero(x > threshold)
print("nResultado de usar nonzero:")
print(indices2)

El resultado de salida es:

原始张量:
tensor([[-1.2345,  0.5678, -0.8901,  1.2345],
        [ 0.3456, -0.6789,  0.9012, -0.1234],
        [-0.5678,  1.2345, -0.3456,  0.7890]])

大于 0 的元素索引:
tensor([[0, 1],
        [0, 3],
        [1, 0],
        [1, 2],
        [2, 1],
        [2, 3]])

使用 nonzero 的结果:
tensor([[0, 1],
        [0, 3],
        [1, 0],
        [1, 2],
        [2, 1],
        [2, 3]])

Ejemplo

import torch

# Tensor 1D
x = torch.tensor([1, 0, 0, 4, 0, 5, 0])

indices = torch.argwhere(x)
print("Índices del tensor 1D:")
print(indices.squeeze())  # Eliminar las dimensiones redundantes

El resultado de salida es:

1D张量的索引:
tensor([0, 3, 5])

Ejemplo

import torch

# Aplicación: encontrar los elementos que cumplen la condición y modificarlos
x = torch.randn(5, 5)

# Encontrar los índices de todos los elementos mayores que 0
pos_indices = torch.argwhere(x > 0)

print("Posiciones mayores que 0 en el tensor original:")
print(pos_indices)

# Usar estos índices para modificar los valores
for idx in pos_indices:
    x[idx[0], idx[1]] = 100

print("nTensor modificado:")
print(x)

Nota:torch.argwhereSítorch.nonzeroalias de, ambas funciones tienen exactamente la misma funcionalidad. Devuelve los índices de los elementos; cada fila corresponde a la posición de un elemento que cumple la condición en el tensor original.

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

Otras extensiones