torch.gather

最新推荐文章于 2023-11-21 11:35:40 发布

kczhan

最新推荐文章于 2023-11-21 11:35:40 发布

阅读量507

点赞数

文章标签： pytorch

本文链接：https://blog.csdn.net/kczhan/article/details/98764885

版权

torch.gather

参数：torch.gather(input, dim, index, out=None)
注意：
· output.shape == index.shape
· index.dtype == torch.LongTensor
· index不能为负数，-1也会报错
分析：对于3维张量
· dim=0, output(i,j,k) = input(index(i,j,k), j, k)
· dim=1, output(i,j,k) = input(i, index(i,j,k), k)
· dim=2, output(i,j,k) = input(i, j, index(i,j,k))

例子1: 二维

In [1]: x = torch.arange(4).reshape(2,2)
In [2]: x
Out[2]: tensor([[0, 1],
    			[2, 3]])     			
In [3]: gather_index = torch.LongTensor([[0, 1],[0, 1]])
In [4]: torch.gather(x, dim=1, index=gather_index)
Out[4]: tensor([[0, 1],
    	 	    [2, 3]])

例子2: 三维

In [1]: x = torch.arange(12).reshape(2,3,2)
In [2]: x
Out[2]: tensor([[[0, 1],
     			 [2, 3],
     		     [4, 5]],
       		    [[6, 7],
        		 [8, 9],
        		 [10, 11]]])	
In [3]: gather_index = gather_index.unsqueeze(-1).expand(x.shape[0],2,x.shape[-1])
In [4]: gather_index
Out[4]:tensor([[[0, 0],
     			[1, 1]],
    		   [[0, 0],
     			[1, 1]]])
In [5]: torch.gather(x, dim=1, index=gather_index)
Out[5]: tensor([[[0, 1],
     			 [2, 3]],
                [[6, 7],
                 [8, 9]]])