參數(shù)說明

以官方說明為例，gather()函數(shù)需要三個(gè)參數(shù)，輸入input，維度dim，以及索引index

input必須為Tensor類型

dim為int類型，代表從哪個(gè)維度進(jìn)行索引

index為LongTensor類型

舉例說明

input=torch.tensor([[1,2,3],[4,5,6]]) #作為輸入

index1=torch.tensor([[0,1,1],[0,1,1]]) #作為索引矩陣

# dim=0時(shí),按列進(jìn)行索引
print (torch.gather(input,dim=0,index=index1))

# dim=1時(shí),按行進(jìn)行索引
print (torch.gather(input,dim=1,index=index1))

結(jié)果如下圖所示：

# 按列進(jìn)行索引
tensor([[1, 5, 6],
      [4, 2, 6]])

# 按行進(jìn)行索引
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)容請搜索服務(wù)器之家以前的文章或繼續(xù)瀏覽下面的相關(guān)文章希望大家以后多多支持服務(wù)器之家！

原文鏈接：https://blog.csdn.net/xiaoliujun1999/article/details/121292061