聽名字就知道這個函式是用來求tensor中某個dim的前k大或者前k小的值以及對應的index。

用法

torch.topk(input,k,dim=None,largest=True,sorted=True,out=None) -> (Tensor,LongTensor)

input：一個tensor資料

k：指明是得到前k個數據以及其index

dim： 指定在哪個維度上排序， 預設是最後一個維度

largest：如果為True，按照大到小排序； 如果為False，按照小到大排序

sorted：返回的結果按照順序返回

out：可預設，不要

topk最常用的場合就是求一個樣本被網路認為前k個最可能屬於的類別。我們就用這個場景為例，說明函式的使用方法。

假設一個，N是樣本數目，一般等於batch size， D是類別數目。我們想知道每個樣本的最可能屬於的那個類別，其實可以用torch.max得到。如果要使用topk，則k應該設定為1。

import torch

pred = torch.randn((4,5))
print(pred)
values,indices = pred.topk(1,dim=1,sorted=True)
print(indices)
# 用max得到的結果，設定keepdim為True，避免降維。因為topk函式返回的index不降維，shape和輸入一致。
_,indices_max = pred.max(dim=1,keepdim=True)

print(indices_max == indices)
# pred
tensor([[-0.1480,-0.9819,-0.3364,0.7912,-0.3263],[-0.8013,-0.9083,0.7973,0.1458,-0.9156],[-0.2334,-0.0142,-0.5493,0.0673,0.8185],[-0.4075,-0.1097,0.8193,-0.2352,-0.9273]])
# indices,shape為 【4,1】,tensor([[3],#【0,0】代表 第一個樣本最可能屬於第一類別
    [2],# 【1,0】代表第二個樣本最可能屬於第二類別
    [4],[2]])
# indices_max等於indices
tensor([[True],[True],[True]])

現在在嘗試一下k=2

import torch

pred = torch.randn((4,indices = pred.topk(2,sorted=True) # k=2
print(indices)
# pred
tensor([[-0.2203,-0.7538,1.8789,0.4451,-0.2526],[-0.0413,0.6366,1.1155,0.3484,0.0395],[ 0.0365,0.5158,1.1067,-0.9276,-0.2124],[ 0.6232,0.9912,-0.8562,0.0148,1.6413]])
# indices
tensor([[2,3],[2,1],[4,1]])

可以發現indices的shape變成了【4,k】，k=2。

其中indices[0] = [2,3]。其意義是說明第一個樣本的前兩個最大概率對應的類別分別是第3類和第4類。

大家可以自行print一下values。可以發現values的shape和indices的shape是一樣的。indices描述了在values中對應的值在pred中的位置。

以上這篇PyTorch中topk函式的用法詳解就是小編分享給大家的全部內容了，希望能給大家一個參考，也希望大家多多支援我們。