python - 仅使用一维索引多维张量

Question

b我有一个形状为： 的 PyTorch 张量torch.Size([10, 10, 51])。我想使用 numpy 数组在维度 d=1（中间一个）的 10 个可能元素之间选择一个元素a = np.array([0,1,2,3,4,5,6,7,8,9])：这只是一个随机的例子。

我想做： b[:,a,:]但这行不通

Answer 1

您的解决方案很可能torch.index_select（文档）

不过，你必须先变成a张量。

a_torch = torch.from_numpy(a)
answer = torch.index_select(b, 1, a_torch)

Answer 2

b在第二个轴上使用索引a应该做：

>>> b = torch.rand(10, 10, 51)
>>> a = np.array([0,1,2,3,4,5,6,7,8,9])

>>> b[:,  a].shape
torch.Size([10, 10, 51])

Answer 3

我在 PyTorch 论坛上找到了解决方案：（https://discuss.pytorch.org/t/how-to-select-specific-vector-in-3d-tensor-beautifully/37724）

x = torch.tensor([[[1, 2, 3],
                   [4, 5, 6],
                   [7, 8, 9]],
                  [[11, 12, 13],
                   [14, 15, 16],
                   [17, 18, 19]]])

idx = torch.tensor([1, 2])
x[torch.arange(x.size(0)), idx]