解析Pytorch中的torch.gather()函數
參數說明
以官方說明為例,gather()函數需要三個參數,輸入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))
結果如下圖所示:
# 按列進行索引 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]])
到此這篇關於Pytorch中的torch.gather()函數的文章就介紹到這瞭,更多相關Pytorch torch.gather()函數內容請搜索WalkonNet以前的文章或繼續瀏覽下面的相關文章希望大傢以後多多支持WalkonNet!
推薦閱讀:
- pytorch人工智能之torch.gather算子用法示例
- Pytorch數據類型與轉換(torch.tensor,torch.FloatTensor)
- PyTorch中Tensor和tensor的區別及說明
- pytorch中index_select()的用法詳解
- pytorch中常用的乘法運算及相關的運算符(@和*)