PyTorch torch.argwhere 関数
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)
# テンソルを作成する
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)
# ブール条件
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)
# ある値より大きい要素を見つける
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テンソル
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)
# 応用:条件を満たす要素を見つけて修正する
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 リファレンスマニュアル