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]])