torch.gather(input, dim, index, out=None) → Tensor

我們直接用例子來說明好了

以三位張量為例

out[i] [j] [k] = tensor[index[i][j][k]] [j] [k] # dim=0out[i] [j] [k] = tensor[i] [index[i][j][k]] [k] # dim=1out[i] [j] [k] = tensor[i] [j] [index[i][j][k]] # dim=3 import torch a = torch.Tensor([[1,2],[3,4]]) torch.gather(a,0,index=torch.LongTensor([[0,0],[1,0]])) ''' tensor([[1., 2.],[3., 2.]]) '''

這個怎么看呢

out[0][0]	a[index[0][0]]? [0]]	a[0][0]=1
out[1][0]	a[index[1][0]]? [0]]	a[1][0] =3
out[0][1]	a[index[0][1]]? [1]]? ?	a[0][1]=2
out[1][1]	a[index[1][1]]? [1]]	a[0][1]=2

還有一種用法在每一行選擇一個的時候比較常用

import torch a = torch.Tensor([[1,2,5,6],[3,4,7,8]]) torch.gather(a,1,torch.LongTensor([[0],[3]])) ''' tensor([[1.],[8.]]) '''

總結

以上是生活随笔為你收集整理的pytorch 笔记：gather 函数的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。