python - torch 聚中维

Question

让 a成为 (n, d, l)张量。让 indices成为 (n, 1)张量，包含索引。我要从a收集在来自 indices 给出的索引的中间维度张量中.因此，生成的张量的形状为 (n, l) .

n = 3
d = 2
l = 3

a = tensor([[[ 0,  1,  2],
             [ 3,  4,  5]],

            [[ 6,  7,  8],
             [ 9, 10, 11]],

            [[12, 13, 14],
             [15, 16, 17]]])

indices = tensor([[0],
                  [1],
                  [0]])

# Shape of result is (n, l)
result = tensor([[ 0,  1,  2],  # a[0, 0, :] since indices[0] == 0

                 [ 9, 10, 11],  # a[1, 1, :] since indices[1] == 1

                 [12, 13, 14]]) # a[2, 0, :] since indices[2] == 0

这确实类似于 a.gather(1, indices) ，但是 gather将无法工作，因为 indices与 a 的形状不同.我如何使用 gather在这种情况下？或者我应该用什么？

Answer 1

您可以手动创建索引。 indices如果张量具有示例数据的形状，则必须将其展平。

a[torch.arange(len(a)),indices.view(-1)]
# equal to a[[0,1,2],[0,1,0]]

出去:

tensor([[ 0,  1,  2],
        [ 9, 10, 11],
        [12, 13, 14]])