解析Pytorch中的torch.gather()函數(shù)
發(fā)布日期:2021-12-11 08:56 | 文章來源:源碼中國
參數(shù)說明
以官方說明為例,gather()函數(shù)需要三個參數(shù),輸入input,維度dim,以及索引index
input必須為Tensor類型
dim為int類型,代表從哪個維度進行索引
index為LongTensor類型
舉例說明
input=torch.tensor([[1,2,3],[4,5,6]]) #作為輸入 index1=torch.tensor([[0,1,1],[0,1,1]]) #作為索引矩陣 # dim=0時,按列進行索引 print (torch.gather(input,dim=0,index=index1)) # dim=1時,按行進行索引 print (torch.gather(input,dim=1,index=index1))
結(jié)果如下圖所示:
# 按列進行索引 tensor([[1, 5, 6], [4, 2, 6]]) # 按行進行索引 tensor([[1, 2, 2], [5, 4, 5]])
畫圖說明


官方文檔
def gather(self, input, dim, index, *args, **kwargs): For a 3-D tensor the output is specified by:: out[i][j][k] = input[index[i][j][k]][j][k] # if dim == 0 out[i][j][k] = input[i][index[i][j][k]][k] # if dim == 1 out[i][j][k] = input[i][j][index[i][j][k]] # if dim == 2 Args: input (Tensor): the source tensor dim (int): the axis along which to index index (LongTensor): the indices of elements to gather Example:: >>> t = torch.tensor([[1, 2], [3, 4]]) >>> torch.gather(t, 1, torch.tensor([[0, 0], [1, 0]])) tensor([[ 1, 1], [ 4, 3]])
到此這篇關(guān)于Pytorch中的torch.gather()函數(shù)的文章就介紹到這了,更多相關(guān)Pytorch torch.gather()函數(shù)內(nèi)容請搜索本站以前的文章或繼續(xù)瀏覽下面的相關(guān)文章希望大家以后多多支持本站!
版權(quán)聲明:本站文章來源標注為YINGSOO的內(nèi)容版權(quán)均為本站所有,歡迎引用、轉(zhuǎn)載,請保持原文完整并注明來源及原文鏈接。禁止復制或仿造本網(wǎng)站,禁止在非maisonbaluchon.cn所屬的服務(wù)器上建立鏡像,否則將依法追究法律責任。本站部分內(nèi)容來源于網(wǎng)友推薦、互聯(lián)網(wǎng)收集整理而來,僅供學習參考,不代表本站立場,如有內(nèi)容涉嫌侵權(quán),請聯(lián)系alex-e#qq.com處理。
相關(guān)文章
關(guān)注官方微信