PyTorch 中的批量 index_fill

Question

我有一个大小为(2, 3)的索引张量：

>>> index = torch.empty(6).random_(0,8).view(2,3)
tensor([[6., 3., 2.],
        [3., 4., 7.]])

还有一个大小为(2, 8)的值张量：

>>> value = torch.zeros(2,8)
tensor([[0., 0., 0., 0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0., 0., 0., 0.]])

我想通过沿dim=-1 的索引将value 中的元素设置为1。** 输出应该是这样的：

>>> output
tensor([[0., 0., 1., 1., 0., 0., 1., 0.],
        [0., 0., 0., 1., 1., 0., 0., 1.]])

我尝试了value[range(2), index] = 1，但它触发了错误。我也试过torch.index_fill，但它不接受批量索引。 torch.scatter 需要创建一个大小为2*8 的额外张量，其中包含1，这会消耗不必要的内存和时间。

Answer 1

您实际上可以通过设置 value (int) 选项而不是 src 选项 (Tensor) 来使用 torch.Tensor.scatter_。

>>> value.scatter_(dim=-1, index=index.long(), value=1)

>>> value
tensor([[0., 0., 1., 1., 0., 0., 1., 0.],
        [0., 0., 0., 1., 1., 0., 0., 1.]])

确保index 的类型是int64。