关于Pytorch中gather的一些使用心得

使用心得"/>

关于Pytorch中gather的一些使用心得

问题场景描述

我在复现Faster-RCNN模型的过程中遇到这样一个问题：
有一个张量，它的形状是 (128, 21, 4)

roi_loc.shape = (128, 21, 4)

与之对应的还有一个label数据

gt_label.shape = (128)

我现在的需求是将label当作第一个张量在dim=1上的索引，将其中的数据拿出来。
具体来说就是，现在有128个样本数据，每个样本中有21个长度为4的向量。label也是128个，每个值代表取出21个向量中的哪一个。

问题的思考

我尝试了很多办法，包括布尔索引，index_select方法等，最后发现都不适用（也有可能我没用好）。最后利用gather API解决了这个问题。

这个API的说明我看了很多遍都没看懂，我相信绝大部分读者也是因为看不懂这个说明才来这儿的。
下面我给出自己的一些理解：

gather的说明

gather所需要的第一个参数是待索引的数据，在我们的问题中 roi_loc就是这个input。第二个参数dim，是你的索引数据要作用在哪个轴上，正如前面所言，我们想索引第二个轴(dim=1).
最难理解的是index，index就是我们想要用来索引的张量，对应的是label。可是label不能直接拿来用，得先做一定的变换，这也就是gather的难点。

我们先从简单的情况来看

input和gather必须在维度上相同，假设数据还是3 * 3，index也是1 * 3的(注意这里是二维的)

此时row至多取值0，col至多取值为2
如果我要对dim=0索引
那么data[0][0] = data[index[0][0]] [0] = data[1][0] = 2
data[0][1] = data[index[0][1]] [1] = data[0][1] = 5
data[0][2] = data[index[0][2]][2] = data[2][2] = 9
上面的过程可以描述为，第一列的元素我想选第二行的，第二列的元素我想选第一行的，第三列的元素我想选第三行的。
可以发现因为index是1 * 3的，所以最后的输出也是31* 3，即输出张量的shape取决于index的shape

以上过程我相信读者好好体悟应该可以理解。

问题的解决

回到我们的问题

roi_loc.shape = (128, 21, 4)，gt_label.shape = (128)

我们想索引dim=1，最后的结果应该是(128, 4)
由上面的说明可以知道，input和index的dimension首先得相同

idx = gt_roi_labels.unsqueeze(-1).unsqueeze(-1)
idx.shape = (128, 1, 1)

又因为我们想要输出的结果得是（128， 4），所以得让idx在最后一个轴上重复4次

idx = idx.repeat_interleave(-1, 4)
idx.shape = (128, 1, 4)

现在就可以利用gather在dim=1上索引了

result = roi_loc.gather(1, idx)
result.shape = (128, 1, 4)

最后将长度为1的轴压缩(本身这个轴的出现是为了满足input和index维度一样的要求)

result = result.squeeze(1)
result.shape(128, 4)