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

末蓝、 2022-09-11 09:25 213阅读 0赞

# 问题场景描述 #

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

roi_loc.shape = (128, 21, 4)

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

gt_label.shape = (128)

![在这里插入图片描述][watermark_type_ZHJvaWRzYW5zZmFsbGJhY2s_shadow_50_text_Q1NETiBA566X5rOV6I-c6bif6aOe6auY6auY_size_20_color_FFFFFF_t_70_g_se_x_16]

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

# 问题的思考 #

我尝试了很多办法，包括布尔索引，index\_select方法等，最后发现都不适用（也有可能我没用好）。最后利用gather API解决了这个问题。  
![在这里插入图片描述][watermark_type_ZHJvaWRzYW5zZmFsbGJhY2s_shadow_50_text_Q1NETiBA566X5rOV6I-c6bif6aOe6auY6auY_size_20_color_FFFFFF_t_70_g_se_x_16 1]  
这个API的说明我看了很多遍都没看懂，我相信绝大部分读者也是因为看不懂这个说明才来这儿的。  
下面我给出自己的一些理解：

# gather的说明 #

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

我们先从简单的情况来看

input和gather必须在维度上相同，假设数据还是3 \* 3，index也是1 \* 3的(注意这里是二维的)  
![在这里插入图片描述][watermark_type_ZHJvaWRzYW5zZmFsbGJhY2s_shadow_50_text_Q1NETiBA566X5rOV6I-c6bif6aOe6auY6auY_size_15_color_FFFFFF_t_70_g_se_x_16]  
此时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  
![在这里插入图片描述][watermark_type_ZHJvaWRzYW5zZmFsbGJhY2s_shadow_50_text_Q1NETiBA566X5rOV6I-c6bif6aOe6auY6auY_size_20_color_FFFFFF_t_70_g_se_x_16 2]

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

# 问题的解决 #

回到我们的问题

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)

[watermark_type_ZHJvaWRzYW5zZmFsbGJhY2s_shadow_50_text_Q1NETiBA566X5rOV6I-c6bif6aOe6auY6auY_size_20_color_FFFFFF_t_70_g_se_x_16]: /images/20220828/7d877bb40c434188b7dd7a298013f0d9.png
[watermark_type_ZHJvaWRzYW5zZmFsbGJhY2s_shadow_50_text_Q1NETiBA566X5rOV6I-c6bif6aOe6auY6auY_size_20_color_FFFFFF_t_70_g_se_x_16 1]: /images/20220828/6e8dba8c0a754cc4bf02a464962c61a0.png
[watermark_type_ZHJvaWRzYW5zZmFsbGJhY2s_shadow_50_text_Q1NETiBA566X5rOV6I-c6bif6aOe6auY6auY_size_15_color_FFFFFF_t_70_g_se_x_16]: /images/20220828/3aabd16c3e05408a837404212f76c283.png
[watermark_type_ZHJvaWRzYW5zZmFsbGJhY2s_shadow_50_text_Q1NETiBA566X5rOV6I-c6bif6aOe6auY6auY_size_20_color_FFFFFF_t_70_g_se_x_16 2]: /images/20220828/0140b297c2444d5a9310543ac4844dfd.png