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

作者:

文章目录

【摘自:Pytorch:torch.nonzero()函数

一维例子:

其中的原理是,a>6得到的是一个元素为True或False的张量,如下图,True会被torch.nonzero()函数认为是非零元素,False会被认为是零元素。

二维例子:

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

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

一维例子:

二维例子:

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

评论

发表回复

您的邮箱地址不会被公开。 必填项已用 * 标注

站内搜索