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

问题场景描述

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

roi_loc.shape = (128, 21, 4)

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

gt_label.shape = (128)

Pytorch深度学习gather一些使用问题解决方案_第1张图片

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

问题的思考

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

Pytorch深度学习gather一些使用问题解决方案_第2张图片

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

gather的说明

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

我们先从简单的情况来看

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

Pytorch深度学习gather一些使用问题解决方案_第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

Pytorch深度学习gather一些使用问题解决方案_第4张图片

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

问题的解决

回到我们的问题

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使用问题的资料请关注脚本之家其它相关文章!

你可能感兴趣的