根据张量元素在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])

这里，指数之间 Offsets[i] 和 Offsets[i+1] 输入数组中。ValuesToAdd[i] 被添加。例如，对于指数 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