如何使用列表作为pytorch的索引答案

【问题标题】：How to use list as index with pytorch如何使用列表作为pytorch的索引
【发布时间】：2020-06-06 16:56:21
【问题描述】：

例如，一个二维张量：

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

还有一个列表l = [0, 1]

如果我执行t[l]，那么它最终会打印t 的第0 行和第一行。

但是如果我想使用l 作为索引呢？我希望使用l 来查找第 0 行和第 1 列的元素。换句话说，我希望得到与t[0, 1] 或t[0][1] 相同的结果。

而且我还想在超过 2d 的维度上使用它。使用长度为n 的l 作为索引来跟踪n 维度张量中的元素。

【问题讨论】：

【解决方案1】：

IIUC 您可以针对给定的场景执行此操作 - t[tuple(l)]

t
tensor([[1, 2, 3],
        [4, 5, 6],
        [7, 8, 9]])
l
[0, 1]

t[tuple(l)]        # equivalent to t[(0,1)] or t[0,1]
tensor(2)

【讨论】：

【解决方案2】：

我刚刚写了一个递归函数来解决这个问题，也许有人有更优雅的解决方案？

def list_as_index(t, l):
    if not l:
        return t
    else:
        return list_as_index(t[l[0]], l[1:])

【讨论】：