torch.gather()的使用举例

参考链接: gather(dim, index)
参考链接: torch.gather(input, dim, index, out=None, sparse_grad=False)

使用说明:

input和index两个参数的形状要满足一定的条件,
即,除了dim参数指定的轴上,其他所有的轴形状必须完全相同.
计算得到的输出结果out的形状和index相同,
out输出结果的值由input给出,具体在input中的位置由
当前out输出结果的位置以及index所决定,
dim轴的位置由index给出,其他轴上的位置由out的当前输出结果位置决定.

代码实验展示:

Microsoft Windows [版本 10.0.18363.1256]
(c) 2019 Microsoft Corporation。保留所有权利。

C:\Users\chenxuqi>conda activate ssd4pytorch1_2_0

(ssd4pytorch1_2_0) C:\Users\chenxuqi>python
Python 3.7.7 (default, May  6 2020, 11:45:54) [MSC v.1916 64 bit (AMD64)] :: Anaconda, Inc. on win32
Type "help", "copyright", "credits" or "license" for more information.
>>> import torch
>>>
>>> torch.manual_seed(seed=20200910)
<torch._C.Generator object at 0x0000019539C5D330>
>>>
>>> t = torch.tensor([[1,2],[3,4]])
>>> t
tensor([[1, 2],
        [3, 4]])
>>> torch.gather(t, 1, torch.tensor([[0,0],[1,0]]))
tensor([[1, 1],
        [4, 3]])
>>> index = torch.tensor([[0,0],[1,0]])
>>> index
tensor([[0, 0],
        [1, 0]])
>>> torch.gather(t, 1, index )
tensor([[1, 1],
        [4, 3]])
>>> torch.gather(t, 0, index)
tensor([[1, 2],
        [3, 2]])
>>>
>>>
>>>