如何在索引处添加到 pytorch 张量？答案

【问题标题】：How to add to pytorch tensor at indices?如何在索引处添加到 pytorch 张量？
【发布时间】：2019-09-20 03:01:45
【问题描述】：

我不得不承认，我对 scatter* 和 index* 操作有点困惑 - 我不确定它们是否完全符合我的要求，这很简单：

给定一些二维张量

z = tensor([[1., 1., 1., 1.],
            [1., 1., 1., 1.],
            [1., 1., 1., 1.]])

还有一个二维索引列表（或张量？）：

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

我想在这些索引处为 z 添加一个标量（并有效地做到这一点）：

znew = z.something_add(inds, 3)
->
znew = tensor([[4., 1., 1., 1.],
               [1., 4., 4., 1.],
               [1., 1., 1., 1.]])

如果必须，我可以将该标量设为任何形状的张量（所有元素 = 3），但我宁愿不...

【问题讨论】：

标签： pytorch

【解决方案1】：

您必须为索引提供两个列表。第一个具有行位置，第二个具有列位置。在您的示例中，它将是：

z[[0, 1, 1], [0, 1, 2]] += 3

【讨论】：

我接受你的，因为它有效而且你是第一个。不过，我仍然有顾虑——这会有效吗？如果 z 和 inds 都是 gpu 驻留张量，这一切都可以在 gpu 上完成吗？似乎张量上应该有一些方法可以做到这一点，并有效地做到这一点......
我认为它应该是高效的，即使你正在传递 CPU 列表。我不确定是否有巨大的索引数组。如果你的张量很大，但你只关心部分组件，试试torch.sparse。

【解决方案2】：

这段代码实现了你想要的：

z_new = z.clone() # copy the tensor
z_new[inds[:, 0], inds[:, 1]] += 3 # modify selected indices of new tensor

在 PyTorch 中，您可以使用另一个张量索引一个张量的每个轴。

【讨论】：