使用2d张量索引3d张量

Question

我有一个3d张量，source形状(bsz x slen1 x nhd)和2d张量，index形状(bsz x slen2)。更具体地说，我有：

source = 32 x 20 x 768
index  = 32 x 16

index张量中的每个值都在[0, 19]之间，source是根据32 x 16 x 768张量的第二个dim的期望向量的索引。

索引后，我期待输出张量的形状，

bsz, _, nhid = source.size()
_, slen = index.size()

source = source.reshape(-1, nhid)
source = source[index.reshape(-1), :]
source = source.reshape(bsz, slen, nhid)

。

目前我这样做：

source = torch.FloatTensor([
    [[ 0.2413, -0.6667,  0.2621],
     [-0.4216,  0.3722, -1.2258],
     [-0.2436, -1.5746, -0.1270],
     [ 1.6962, -1.3637,  0.8820],
     [ 0.3490, -0.0198,  0.7928]],

    [[-0.0973,  2.3106, -1.8358],
     [-1.9674,  0.5381,  0.2406],
     [ 3.0731,  0.3826, -0.7279],
     [-0.6262,  0.3478, -0.5112],
     [-0.4147, -1.8988, -0.0092]]
     ])

index = torch.LongTensor([[0, 1, 2, 3], 
                          [1, 2, 3, 4]])

所以，我正在将3d源张量转换为2d张量和2d索引张量到1d张量，然后执行索引。它是否正确？

有没有更好的方法呢？

更新

我检查了我的代码没有给出预期的结果。为了解释我想要的，我提供以下代码片段。

torch.FloatTensor([
    [[ 0.2413, -0.6667,  0.2621],
     [-0.4216,  0.3722, -1.2258],
     [-0.2436, -1.5746, -0.1270],
     [ 1.6962, -1.3637,  0.8820]],

    [[-1.9674,  0.5381,  0.2406],
     [ 3.0731,  0.3826, -0.7279],
     [-0.6262,  0.3478, -0.5112],
     [-0.4147, -1.8988, -0.0092]]
     ])

我希望输出张量为：

source[torch.arange(source.shape[0]).unsqueeze(-1), index]

Answer 1

更新：

torch.arange(source.shape[0]).unsqueeze(-1)

请注意，

tensor([[0],
        [1]])  # 2 x 1

给出：

index

和

tensor([[0, 1, 2, 3],
        [1, 2, 3, 4]])  # 2 x 4

是：

arange

index对批量维度进行索引，而slen1同时对unsqueeze维度进行索引。 x 1调用将额外的arange维度添加到

index = torch.LongTensor([[0, 1, 2, 3], [1, 2, 3, 4]])
offset = torch.arange(0, source.size(0) * source.size(1), source.size(1))
index = index + offset.unsqueeze(1)

source = source.reshape(-1, source.shape[-1])[index]

结果中，以便两者可以一起广播。

Answer 2

我已经解决了这个问题。所以，我实际上需要定义偏移量。以下代码适用于我。

qazxswpoi