torch.nonzero()和torch.index_select()，筛选张量中符合某种条件的元素

1. torch.nonzero()的定义

【摘自：Pytorch：torch.nonzero()函数】

 

2. torch.nonzero()用来筛选张量中符合某个条件的元素，得到这些元素的索引

一维例子：

注意得到的b是二维张量。上述代码的原理是，a>6得到的是一个元素为True或False的张量，如下图，True会被torch.nonzero()函数认为是非零元素，False会被认为是零元素。

 

二维例子：

 

3. torch.index_select()的定义和示例

torch.index_select(input, dim, index) 函数返回的是沿着输入张量的指定维度的指定索引号进行索引的张量子集，函数参数有：

• input(Tensor) - 需要进行索引操作的输入张量；
• dim(int) - 需要对输入张量进行索引的维度；
• index(LongTensor) - 包含索引号的 1D 张量；

一维例子：

 

 

二维例子：

 

4. torch.nonzero()和torch.index_select()结合使用

结合使用torch.nonzero()和torch.index_select()，可以选出符合某种条件的元素。下面的例子是从一维张量a中选出大于6的元素：