Python PyTorch bucketize用法及代码示例

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

用法:

torch.bucketize(input, boundaries, *, out_int32=False, right=False, out=None) → Tensor

参数：

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

• boundaries(Tensor) -一维张量，必须包含一个单调递增的序列。

关键字参数：

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

• right(bool,可选的) -如果为 False，则返回找到的第一个合适的位置。如果为 True，则返回最后一个此类索引。如果没有找到合适的索引，则返回 0 表示非数值(例如 nan、inf)或 boundaries 的大小(通过最后一个索引)。换句话说，如果为 False，则从 boundaries 获取 input 中每个值的下限索引。如果为 True，则改为获取上限索引。默认值为假。

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

返回 input 中每个值所属的桶的索引，其中桶的边界由 boundaries 设置。返回一个与 input 大小相同的新张量。如果 right 为 False(默认)，则左边界是闭合的。更正式地说，返回的索引满足以下规则：

right	返回的索引满足
False	boundaries[i-1] < input[m][n]...[l][x] <= boundaries[i]
True	boundaries[i-1] <= input[m][n]...[l][x] < boundaries[i]

例子：

>>> boundaries = torch.tensor([1, 3, 5, 7, 9])
>>> boundaries
tensor([1, 3, 5, 7, 9])
>>> v = torch.tensor([[3, 6, 9], [3, 6, 9]])
>>> v
tensor([[3, 6, 9],
        [3, 6, 9]])
>>> torch.bucketize(v, boundaries)
tensor([[1, 3, 4],
        [1, 3, 4]])
>>> torch.bucketize(v, boundaries, right=True)
tensor([[2, 3, 5],
        [2, 3, 5]])