问题场景描述
我在复现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)
以上就是Pytorch深度学习gather一些使用问题解决方案的详细内容，更多关于Pytorch学习gather使用问题的资料请关注服务器之家其它相关文章！
原文链接：https://blog.csdn.net/qq_43152622/article/details/120304409