Pytorch 可微分条件（基于索引）和

Question

我有一个 idx 数组，如 [0, 1, 0, 2, 3, 1] 和另一个二维数组 data，如下所示：

[[0,  1,  2],
 [3,  4,  5],
 [6,  7,  8],
 [9,  10, 11],
 [12, 13, 14],
 [15, 16, 17]]

我希望我的输出是 4x3，其中 4 是 idx 的最大值，3 是特征大小 (data.shape[1])，并且在输出中每个元素是特征的总和，在 @ 987654329@。那么本例中的输出将是：

[[6,  8,  10],
 [18, 20, 22],
 [9,  10, 11],
 [12, 13, 14]]

我可以通过迭代 range(3) 并在数据上创建一个掩码并总结它们来做到这一点，但它是不可区分的（我想）。 Pytorch 中是否有任何功能用于此目的？类似scatter()。

更新：我似乎正在寻找一个名为 scatter sum 的东西，它在 this 存储库中实现。

Answer 1

您正在寻找index_add_:

import torch

x = torch.tensor([[ 0.,  1.,  2.],
        [ 3.,  4.,  5.],
        [ 6.,  7.,  8.],
        [ 9., 10., 11.],
        [12., 13., 14.],
        [15., 16., 17.]], dtype=torch.float)
idx = torch.tensor([0, 1, 0, 2, 3, 1], dtype=torch.long)  # note the dtype here, must be "long"
# init the sums to zero
y = torch.zeros((idx.max()+1, x.shape[1]), dtype=x.dtype)

# do the magic
y.index_add_(0, idx, x)

提供所需的输出

tensor([[ 6.,  8., 10.],
        [18., 20., 22.],
        [ 9., 10., 11.],
        [12., 13., 14.]])