Pytorch中的torch.gather()函數(shù)怎么用

這篇文章將為大家詳細(xì)講解有關(guān)Pytorch中的torch.gather()函數(shù)怎么用，小編覺(jué)得挺實(shí)用的，因此分享給大家做個(gè)參考，希望大家閱讀完這篇文章后可以有所收獲。

參數(shù)說(shuō)明

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

input必須為T(mén)ensor類型

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

index為L(zhǎng)ongTensor類型

舉例說(shuō)明

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]])

畫(huà)圖說(shuō)明 

官方文檔

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ù)怎么用”這篇文章就分享到這里了，希望以上內(nèi)容可以對(duì)大家有一定的幫助，使各位可以學(xué)到更多知識(shí)，如果覺(jué)得文章不錯(cuò)，請(qǐng)把它分享出去讓更多的人看到。