获取特定维度的变量索引值答案

【问题标题】：Get value of variable index in particular dimension获取特定维度的变量索引值
【发布时间】：2020-06-21 02:15:37
【问题描述】：

如果我有一个张量是

value = torch.tensor([

    [[0, 0, 0], [1, 1, 1]],
    [[2, 2, 2], [3, 3, 3]],
])

基本上是形状(2,2,3)。

现在说如果我有index = [1, 0]，这意味着我想接受：

# row 1 of [[0, 0, 0], [1, 1, 1]], giving me: [1, 1, 1]
# row 0 of [[2, 2, 2], [3, 3, 3]], giving me: [2, 2, 2]

这样最后的输出：

output = torch.tensor([[1, 1, 1], [2, 2, 2]])

有没有一种矢量化的方式来实现这一点？

【问题讨论】：

标签： python numpy pytorch

【解决方案1】：

您可以使用高级索引。
我找不到关于这个的好的 pytorch 文档，但我相信它和 numpy 一样工作，所以这里是 numpy 的 document about indexing。

import torch

value = torch.tensor([

    [[0, 0, 0], [1, 1, 1]],
    [[2, 2, 2], [3, 3, 3]],
])

index = [1, 0]
i = range(0,2)

result = value[i, index]
# same as result = value[i, index, :] 

print(result)

【讨论】：