pytorch的topk()函数

2022-09-02 21:57:03 浏览数 (2)

pytorch.topk()用于返回Tensor中的前k个元素以及元素对应的索引值。例:

代码语言:javascript复制
import torch

item=torch.IntTensor([1,2,4,7,3,2])
value,indices=torch.topk(item,3)
print("value:",value)
print("indices:",indices)

输出结果为:

 其中:value中存储的是对应的top3的元素,并按照从大到小的取值方式进行存储

indices中存储的是value中top3元素在原Tensor中的索引值

0 人点赞