Pytorch深度学习gather一些使用问题解决方案 Pytorch深度学习gather一些使用问

问题场景描述
问题的思考
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