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

Question

例如，一个二维张量：>>> t = torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]])还有一个清单l = [0, 1]如果我执行t[l]，那么它最终会打印第 0 行和第 1 行t。但是如果我想l用作索引怎么办？我希望用于l查找第 0 行和第 1 列的元素。换句话说，我希望得到与t[0, 1]or相同的结果t[0][1]。我也想在二维以上使用它。使用lwith lengthn作为索引来跟踪n维度张量中的元素。

Answer 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)

Answer 2

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

def list_as_index(t, l):

    if not l:

        return t

    else:

        return list_as_index(t[l[0]], l[1:])