一、特定元素的位置索引
获取矩阵中特定元素的位置索引,可以通过nonzero()结合判断条件来获取。nonzero()函数可以返回非0元素的位置。下面举一下一维以及二维的相关例子。
(1)一维矩阵的nonzero()
>>> import torch
>>> a = torch.tensor([1,0,2,30,0,1])
>>> a.nonzero()
tensor([[0],
[2],
[3],
[5]])
>>> a.nonzero().size()
torch.Size([4, 1])
(2)二维矩阵的nonzero()
>>> b = torch.tensor([[1,0,2],[0,2,1]])
>>> b
tensor([[1, 0, 2],
[0, 2, 1]])
>>> b.nonzero()
tensor([[0, 0],
[0, 2],
[1, 1],
[1, 2]])
>>> b.nonzero().size()
torch.Size([4, 2])
(3)结合判断条件的nonzero(),可以获取指定元素的位置索引
下面例子是获取二维矩阵c中2对应的索引位置。
>>> c = torch.tensor([[1,0,2],[0,2,1]])
>>> c == torch.tensor(2)
tensor([[False, False, True],
[False, True, False]])
>>> (c == torch.tensor(2)).nonzero()
tensor([[0, 2],
[1, 1]])
>>> (c == torch.tensor(2