Pytorch中topk的

 

Pytorch 中topk的用法¶

 

In [14]:

import torch
from IPython.core.interactiveshell import InteractiveShell
InteractiveShell.ast_node_interactivity = "all"a = torch.randint(0,10,[2,2,3])
a
a.topk(2,-1)

 

Out[14]:

tensor([[[4, 1, 5],[9, 2, 4]],[[1, 1, 1],[5, 1, 4]]])

Out[14]:

torch.return_types.topk(
values=tensor([[[5, 4],[9, 4]],[[1, 1],[5, 4]]]),
indices=tensor([[[2, 0],[0, 2]],[[0, 1],[0, 2]]]))

 

 

第一个返回结果为：在指定维度上选出topk个元素，并且依次从大到小排序，输出的形状为 (bach_size,seq_len,top_k) 第二个返回结果为： topk最大的元素索引，其形状为 (batch,seq_len,top_k)