根据它在 pytorch 中的位置向张量元素添加值

Question

在 pytorch 中，我想根据元素的位置向张量中的元素添加值。例如考虑，

Input = torch.tensor([1,2,3,4,5,6,7,8,9,0,1,2,3,4,5,6,7,8,9,0])

在输入数组的几个偏移量之间，Offsets = [0,5,10,15,20]，我想添加不同的值，ValuesToAdd = [10,100,1000,10000]

我希望输出是

Output = torch.tensor([11,12,13,14,15,106,107,108,109,100,1001,1002,1003,1004,1005,10006,10007,10008,10009,10000])

此处，在 Input 数组的索引 Offsets[i] 和 Offsets[i+1] 之间，添加了 ValuesToAdd[i]。例如，对于 Input 数组中的索引 10,11,12,13 和 14（Offsets[2] = 10 到 Offsets[3]=15），添加 1000（ValuesToAdd[2]）。

我怎样才能做到这一点？我正在寻找一种更有效的方法，而不是遍历 Offsets 数组。

Answer 1

您可以使用torch.repeat_interleave

Offsets = torch.tensor(Offsets)
shifts = Offsets[1:] - Offsets[:-1]
output = Input.clone()
output[Offsets[0]:Offsets[-1]] += torch.tensor(ValuesToAdd).repeat_interleave(shifts)
print(torch.all(output == Output))
# True

根据它在 pytorch 中的位置向张量元素添加值

Add values to a tensor element based on its position in pytorch

python

numpy

add

pytorch

tensor