Функция PyTorch torch.argwhere

Pytorch torch 参考手册Справочное руководство PyTorch torch

torch.argwhere— это функция в PyTorch, используемая для возврата индексов элементов, удовлетворяющих условию. Она возвращает индексы элементов, значение которых равно True (ненулевых) во входном тензоре.

Определение функции

torch.argwhere(input)

Parameters:

  • input(Tensor): Входной тензор.

Возвращаемое значение:

  • torch.Tensor: Возвращает двумерный тензор, каждая строка — индекс элемента, удовлетворяющего условию.

Примеры использования

Пример

import torch

# Создайте тензор
x = torch.tensor([[1, 0, 2],
                  [0, 3, 0],
                  [4, 5, 0]])

print("Исходный тензор:")
print(x)

# Возвращаем индексы ненулевых элементов
indices = torch.argwhere(x)

print("nИндексы ненулевых элементов:")
print(indices)

Результат вывода:

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

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

Пример

import torch

# Булево условие
x = torch.tensor([[True, False, True],
                  [False, True, False],
                  [True, True, False]])

print("Булев тензор:")
print(x)

indices = torch.argwhere(x)
print("nИндексы значений True:")
print(indices)

Результат вывода:

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

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

Пример

import torch

# Найти элементы, превышающие заданное значение
x = torch.randn(3, 4)
threshold = 0

print("Исходный тензор:")
print(x)

# Найти индексы элементов, превышающих threshold
indices = torch.argwhere(x > threshold)

print(f"nИндексы элементов, превышающих {threshold}:")
print(indices)

# Вы также можете использовать функцию nonzero с тем же эффектом
indices2 = torch.nonzero(x > threshold)
print("nРезультат с использованием nonzero:")
print(indices2)

Результат вывода:

原始张量:
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]])

Пример

import torch

# 1D тензор
x = torch.tensor([1, 0, 0, 4, 0, 5, 0])

indices = torch.argwhere(x)
print("Индексы 1D тензора:")
print(indices.squeeze())  # Удалить лишнее измерение

Результат вывода:

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

Пример

import torch

# Применение: найти элементы, удовлетворяющие условию, и изменить их
x = torch.randn(5, 5)

# Найдите индексы всех элементов, больших 0
pos_indices = torch.argwhere(x > 0)

print("Позиции в исходном тензоре больше 0:")
print(pos_indices)

# Используйте эти индексы для изменения значений
for idx in pos_indices:
    x[idx[0], idx[1]] = 100

print("nИзменённый тензор:")
print(x)

Примечание:torch.argwhereДаtorch.nonzeroявляется псевдонимом, они имеют абсолютно одинаковую функциональность. Он возвращает индексы элементов, каждая строка соответствует позиции элемента, удовлетворяющего условию в исходном тензоре.

Pytorch torch 参考手册Справочное руководство PyTorch torch

Другие расширения