Pytorch 从最后一个张量维度中选择值，其索引来自另一个维度较小的 Tenor

Question

我有一个张量 a具有三个维度。第一个维度对应于小批量大小，第二个对应于序列长度，第三个对应于特征维度。例如。，

>>> a = torch.arange(1, 13, dtype=torch.float).view(2,2,3)  # Consider the values of a to be random
>>> a
tensor([[[ 1.,  2.,  3.],
         [ 4.,  5.,  6.]],

        [[ 7.,  8.,  9.],
         [10., 11., 12.]]])

我有第二个二维张量。它的第一个维度对应于小批量大小，第二个维度对应于序列长度。它包含 a 的第三维索引范围内的值. a第三个维度的大小为 3，所以 b可以包含值 0、1 或 2。例如，

>>> b = torch.LongTensor([[0, 2],[1,0]])
>>> b
tensor([[0, 2],
        [1, 0]])

我想获得一个张量 c形状为 b并包含 a 的所有值 b 所引用的.
在上面的场景中，我想要:

c = torch.empty(2,2)
c[0,0] = a[0, 0, b[0,0]]
c[1,0] = a[1, 0, b[1,0]]
c[0,1] = a[0, 1, b[0,1]]
c[1,1] = a[1, 1, b[1,1]]

>>> c
tensor([[ 1.,  5.],
        [ 8., 10.]])

如何创建张量 c快速地？此外，我还希望 c 是可微的(能够使用 .backprob() )。我对 pytorch 不太熟悉，所以我不确定是否存在可区分的版本。

作为替代方案，而不是 c具有与 b 相同的形状我也可以使用 c与 a 的形状相同，只有零，但在 b 引用的位置那些。然后我可以乘 a和 c得到一个可微分的张量。

如下:

c = torch.zeros(2,2,3, dtype=torch.float)
c[0,0,b[0,0]] = 1
c[1,0,b[1,0]] = 1
c[0,1,b[0,1]] = 1
c[1,1,b[1,1]] = 1

>>> a*c
tensor([[[ 1.,  0.,  0.],
         [ 0.,  5.,  0.]],

        [[ 0.,  8.,  0.],
         [10.,  0.,  0.]]])

Answer 1

让我们先声明必要的变量:(注意 requires_grad 的初始化中的 a，我们将使用它来确保可微性)

a = torch.arange(1,13,dtype=torch.float32,requires_grad=True).reshape(2,2,3)
b = torch.LongTensor([[0, 2],[1,0]])

让我们 reshape a和 Squash 小批量和序列维度:

temp = a.reshape(-1,3)

所以 temp现在看起来像:

tensor([[ 1.,  2.,  3.],
    [ 4.,  5.,  6.],
    [ 7.,  8.,  9.],
    [10., 11., 12.]], grad_fn=<AsStridedBackward>)

现在注意 b 的每个值可用于 temp的每一行获得所需的输出。现在我们这样做:

c = temp[range(len(temp )),b.view(-1)].view(b.size())

注意我们如何索引 temp , range(len(temp ))选择每一行和一维 b即 b.view(-1)获取对应的列。最后 .view(b.size())使这个数组的大小与 b 相同.

如果我们打印 c现在:

tensor([[ 1.,  6.],
    [ 8., 10.]], grad_fn=<ViewBackward>)

grad_fn=..的存在|显示 c需要梯度，即它是可微的。