PyTorch torch.argwhere 関数


Pytorch torch 参考手册Pytorch torch リファレンスマニュアル

torch.argwherePyTorch において、条件を満たす要素のインデックスを返す関数です。入力テンソル内の値が True(非ゼロ)である要素のインデックスを返します。

関数定義

torch.argwhere(input)

パラメータ:

  • input(Tensor): 入力テンソル。

戻り値:

  • torch.Tensor: 2次元テンソルを返します。各行は条件を満たす要素のインデックスです。

使用例

例

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("nTrue値のインデックス:")
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("nnonzero を使用した結果:")
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 リファレンスマニュアル

その他の拡張