Python PyTorch searchsorted用法及代码示例

本文简要介绍python语言中 torch.searchsorted 的用法。

用法:

torch.searchsorted(sorted_sequence, values, *, out_int32=False, right=False, out=None) → Tensor

参数：

• sorted_sequence(Tensor) -N-D 或一维张量，包含单调递增的序列最里面维度。

• values(Tensor或者标量) -N-D 张量或包含搜索值的标量。

关键字参数：

• out_int32(bool,可选的) -指示输出数据类型。如果为 True，则为 torch.int32，否则为 torch.int64。默认值为 False，即默认输出数据类型为 torch.int64。

• right(bool,可选的) -如果为 False，则返回找到的第一个合适的位置。如果为 True，则返回最后一个此类索引。如果没有找到合适的索引，则返回 0 表示非数值(例如 nan、inf)或最里面内尺寸sorted_sequence(一个通过的最后一个索引最里面方面)。换句话说，如果为 False，则获取中每个值的下限索引values在相应的最里面的维度sorted_sequence.如果为 True，则改为获取上限索引。默认值为假。

• out(Tensor,可选的) -输出张量，必须与values(如果提供)大小相同。

从 中查找索引最里面的维度sorted_sequence这样，如果相应的值在values在索引之前插入，相应的顺序最里面内尺寸sorted_sequence将被保留。返回一个大小相同的新张量values.如果right为 False (默认)，则左边界sorted_sequence已经关闭。更正式地说，返回的索引满足以下规则：

sorted_sequence	right	返回的索引满足
一维	False	sorted_sequence[i-1] < values[m][n]...[l][x] <= sorted_sequence[i]
一维	True	sorted_sequence[i-1] <= values[m][n]...[l][x] < sorted_sequence[i]
N-D	False	sorted_sequence[m][n]...[l][i-1] < values[m][n]...[l][x] <= sorted_sequence[m][n]...[l][i]
N-D	True	sorted_sequence[m][n]...[l][i-1] <= values[m][n]...[l][x] < sorted_sequence[m][n]...[l][i]

注意

如果您的用例始终是一维排序序列，则首选 torch.bucketize() ，因为它具有较少的维度检查，从而导致性能稍好。

例子：

>>> sorted_sequence = torch.tensor([[1, 3, 5, 7, 9], [2, 4, 6, 8, 10]])
>>> sorted_sequence
tensor([[ 1,  3,  5,  7,  9],
        [ 2,  4,  6,  8, 10]])
>>> values = torch.tensor([[3, 6, 9], [3, 6, 9]])
>>> values
tensor([[3, 6, 9],
        [3, 6, 9]])
>>> torch.searchsorted(sorted_sequence, values)
tensor([[1, 3, 4],
        [1, 2, 4]])
>>> torch.searchsorted(sorted_sequence, values, right=True)
tensor([[2, 3, 5],
        [1, 3, 4]])

>>> sorted_sequence_1d = torch.tensor([1, 3, 5, 7, 9])
>>> sorted_sequence_1d
tensor([1, 3, 5, 7, 9])
>>> torch.searchsorted(sorted_sequence_1d, values)
tensor([[1, 3, 4],
        [1, 3, 4]])