Python PyTorch topk用法及代碼示例

本文簡要介紹python語言中 torch.topk 的用法。

用法:

torch.topk(input, k, dim=None, largest=True, sorted=True, *, out=None)

參數：

• input(Tensor) -輸入張量。

• k(int) -“top-k”中的k

• dim(int,可選的) -要排序的維度

• largest(bool,可選的) -控製是否返回最大或最小元素

• sorted(bool,可選的) -控製是否按排序順序返回元素

關鍵字參數：

out(tuple,可選的) -(Tensor, LongTensor) 的輸出元組，可以選擇用作輸出緩衝區

返回給定 input 張量沿給定維度的 k 最大元素。

如果沒有給出dim，則選擇input的最後一個維度。

如果largest 是False，則返回k 最小元素。

返回 (values, indices) 的命名元組，其中 indices 是原始 input 張量中元素的索引。

布爾選項 sorted if True 將確保返回的 k 元素本身是排序的

例子：

>>> x = torch.arange(1., 6.)
>>> x
tensor([ 1.,  2.,  3.,  4.,  5.])
>>> torch.topk(x, 3)
torch.return_types.topk(values=tensor([5., 4., 3.]), indices=tensor([4, 3, 2]))