python - 什么是批量收集的聪明方法？

Question

我有两个矩阵，A 和 B，形状为 (n, m, k) 和 (n, m ) 分别。 n 是批大小，m 是批中的数据量，k 是特征大小。

B 的每个元素都是一个小于m 的索引(特别是B = torch.randint(high=m, shape=(n,m)) )。

我想以更智能的方式实现[A[i][B[i]] for i in range(n)]。

在 pytorch 中有没有更好的方法来实现这个而不用做循环？

Answer 1

你可以使用

a[torch.arange(n)[:, None], b]

一个例子:

>>> n, m, k = 3, 2, 5
>>> a = torch.arange(30).view(n, m, k)
>>> b = torch.randint(high=m, size=(n,m))

# first indexer (of shape (n, 1))
>>> torch.arange(n)[:, None]

tensor([[0],
        [1],
        [2]])

# second indexer
>>> b

tensor([[1, 0],
        [0, 1],
        [1, 1]])

索引器分别具有 (3, 1) 和 (3, 2) 的形状，因此它们将被广播到 (3, 2) 有效地拥有

tensor([[0, 0],
        [1, 1],
        [2, 2]])

和

tensor([[1, 0],
        [0, 1],
        [1, 1]])

表示:对于第一行，取第一个 (k,) 数组并放入结果，然后取第 0 个 (k,) 数组并放入结果。这将在输出中填充 (m, k) 数组，每一行重复 n 次，

得到

>>> a[torch.arange(n)[:, None], b]

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

        [[10, 11, 12, 13, 14],
         [15, 16, 17, 18, 19]],

        [[25, 26, 27, 28, 29],
         [25, 26, 27, 28, 29]]])

与列表理解比较:

>>> [a[i][b[i]] for i in range(n)]

[tensor([[5, 6, 7, 8, 9],
         [0, 1, 2, 3, 4]]),
 tensor([[10, 11, 12, 13, 14],
         [15, 16, 17, 18, 19]]),
 tensor([[25, 26, 27, 28, 29],
         [25, 26, 27, 28, 29]])]