使用 index_select 通过另一个 PyTorch 张量索引一个 PyTorch 张量答案

【问题标题】：Indexing one PyTorch tensor by another using index_select使用 index_select 通过另一个 PyTorch 张量索引一个 PyTorch 张量
【发布时间】：2017-10-23 02:22:40
【问题描述】：

我有一个 3 x 3 PyTorch LongTensor，看起来像这样：

A = 
    [0, 0, 0]
    [1, 2, 2]
    [1, 2, 3]

我想让我们用它来索引一个像这样的 4 x 2 FloatTensor：

B = 
    [0.4, 0.5]
    [1.2, 1.4]
    [0.8, 1.9]
    [2.4, 2.9]

我的预期输出是下面的 2 x 3 x 3 FloatTensor：

C[0,:,:] = 
    [0.4, 0.4, 0.4]
    [1.2, 0.8, 0.8]
    [1.2, 0.8, 2.4]

C[1,:,:] =
    [0.5, 0.5, 0.5]
    [1.4, 1.9, 1.9]
    [1.4, 1.9, 2.9]

换句话说，矩阵A是索引和广播矩阵B。 A是B的索引矩阵，所以这个操作本质上是一个索引操作。

如何使用torch.index_select() 函数来做到这一点？如果解决方案涉及添加或置换维度，那很好。

【问题讨论】：

标签： matrix indexing pytorch

【解决方案1】：

使用index_select() 要求索引值在向量而不是张量中。但只要格式正确，该函数就会为您处理广播。我相信由于广播，必须做的最后一件事是重塑输出。

能成功完成这个操作的单线是

torch.index_select(B, 0, A.view(-1)).view(3,-1,2).permute(2,0,1)

A.view(-1) 向量化索引矩阵。

__.view(3,-1,2) 重新调整为索引矩阵的形状，但考虑了大小为 2 的新额外维度（因为我正在索引 N x 2 矩阵）。

最后，__.permute(2,0,1) 重塑矩阵，以便输出在单独的通道（而不是每一列）中查看 B 的每个维度。

【讨论】：