Pytorch中torch.nonzero()函数解析

一. torch.nonzero()函数解析

1. 官网链接

torch.nonzero()，如下图所示：

2. torch.nonzero()函数解析

2.1. 输入是一维张量，返回一个包含输入 input 中非零元素索引的张量，输出张量中的每行包含 input 中非零元素的索引，输出是二维张量torch.size(z,1)， z 是输入张量 input 中所有非零元素的个数。
2.2. 输入是n维张量，如果输入 input 有 n 维,则输出的索引张量的size为torch.size(z,n) , 这里 z 是输入张量 input 中所有非零元素的个数。
无论输入是几维，输出张量都是两维，每行代表输入张量中非零元素的索引位置(在所有维度上面的位置)。
返回 input中非零元素的索引下标，n维input 中的元素的索引有n 个维度的索引下标。

3. 代码举例

3.1 输入一维张量torch.Size([5])，返回的是输入张量不为零的值在一维向量中的索引位置，输出二维张量torch.Size([4, 1])，4代表输入张量中不为零的个数,1代表在一维张量中的索引位置。

x = torch.tensor([1, 1, 1, 0, 1])
y = torch.nonzero(x)
x,x.shape,y,y.shape

输出结果如下：
(tensor([1, 1, 1, 0, 1]),
 torch.Size([5]),
 tensor([[0],
         [1],
         [2],
         [4]]),
 torch.Size([4, 1]))

3.2 输入二维张量torch.Size([4, 4])，返回的是不为零的值在二维向量中的坐标，例如[0, 0]、[0,2]代表的是二维向量第一行第一列，第一行第三列不为零，输出二维张量torch.Size([6, 2])，6代表输入张量中不为零的个数，2代表在二维张量中的索引位置。
input是2维的，一共有6个非0元素，所以输出是一个torch.Size([6, 2])的张量，表示每个非0元素的索引。读法是从左往右，比如输出张量的第0行[0,0]，表示的就是input的第0行的第0个元素是非0元素，同理，输出张量的第1行[0,2]，表示的就是input的第1行的第2个元素是非0元素等等。

x = torch.tensor([[0.6, 0.0, 0.9, 0.0],
                  [0.0, 0.4, 0.0, 0.0],                            
                  [0.0, 0.0, 1.2, 0.0],
                  [0.0, 0.7, 0.0,-0.4]])
y = torch.nonzero(x)
x,x.shape,y,y.shape

输出结果如下：
(tensor([[ 0.6000,  0.0000,  0.9000,  0.0000],
         [ 0.0000,  0.4000,  0.0000,  0.0000],
         [ 0.0000,  0.0000,  1.2000,  0.0000],
         [ 0.0000,  0.7000,  0.0000, -0.4000]]),
 torch.Size([4, 4]),
 tensor([[0, 0],
         [0, 2],
         [1, 1],
         [2, 2],
         [3, 1],
         [3, 3]]),
 torch.Size([6, 2]))

3.3. 输入三维张量torch.Size([2, 3, 4])，返回的是不为零的值在三维向量中的坐标，将下面例子看作图像，那么图像维度大小[2,3,4]对应于[channels, width, height]，例如[0, 0, 3]代表的是第一个通道第一行第一列不为零，输出二维张量torch.Size([13, 3])，13代表输入张量中不为零的个数，3代表在三维张量中的索引位置。
输出张量按行依次从左往右读，第0个通道第0行第0列的元素非0，第0个通道第1行第0列的元素非0，……，第1个通道第2行第2列的元素非0，第1个通道第2行第2列的元素非0。

x = torch.randn(2,3,4)
y = (x>0.1)
z = torch.nonzero(y)
x,x.shape,y,y.shape,z,z.shape

输出结果如下：
(tensor([[[ 0.3326, -0.9972, -0.4871, -1.3885],
          [ 0.4679, -1.7913,  2.0604,  0.3150],
          [-0.6156, -0.5204,  0.2902, -0.0780]],
 
         [[ 1.2206, -0.7150, -0.1662,  0.5120],
          [ 0.2907,  0.1285,  0.8520, -1.2698],
          [ 0.5176, -0.3800,  0.4408,  0.5073]]]),
 torch.Size([2, 3, 4]),
 tensor([[[ True, False, False, False],
          [ True, False,  True,  True],
          [False, False,  True, False]],
 
         [[ True, False, False,  True],
          [ True,  True,  True, False],
          [ True, False,  True,  True]]]),
 torch.Size([2, 3, 4]),
 tensor([[0, 0, 0],
         [0, 1, 0],
         [0, 1, 2],
         [0, 1, 3],
         [0, 2, 2],
         [1, 0, 0],
         [1, 0, 3],
         [1, 1, 0],
         [1, 1, 1],
         [1, 1, 2],
         [1, 2, 0],
         [1, 2, 2],
         [1, 2, 3]]),
 torch.Size([13, 3]))