PyTorch中topk函式的用法詳解
阿新 • • 發佈:2020-01-09
聽名字就知道這個函式是用來求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函式的用法詳解就是小編分享給大家的全部內容了,希望能給大家一個參考,也希望大家多多支援我們。