使用索引张量选择张量的第二个暗淡答案

【问题标题】：Selecting second dim of tensor using an index tensor使用索引张量选择张量的第二个暗淡
【发布时间】：2019-02-21 05:08:36
【问题描述】：

我有一个二维张量和一个索引张量。 2D 张量有一个批量维度和一个具有 3 个值的维度。我有一个索引张量，它恰好选择 3 个值中的 1 个元素。生成仅包含索引张量中元素的切片的“最佳”方法是什么？

t = torch.tensor([[1,2,3], [4,5,6], [7,8,9]])
t = tensor([[1, 2, 3],
        [4, 5, 6],
        [7, 8, 9]])

i = torch.tensor([0,0,1], dtype=torch.int64)
tensor([0, 0, 1])

预期输出...

tensor([1, 4, 8])

【问题讨论】：

标签： pytorch

【解决方案1】：

答案示例如下。

import torch

t = torch.tensor([[1,2,3], [4,5,6], [7,8,9]])
col_i = [0, 0, 1]
row_i = range(3)
print(t[row_i, col_i])
# tensor([1, 4, 8])

【讨论】：