我有两个张量,第一个包含浮点数,第二个包含 0 和 1。我想根据第二个张量对第一个张量求和。更具体地说,我想在两个 0 的出现之间求和。例如,考虑
a = tensor([1., 1., 1., 1., 1., 1., 1., 1., 1., 1.])
b = tensor([0., 1., 1., 1., 0., 1., 1., 1., 1., 0.])
我想要一些接收两个张量并返回的矢量化(最好)操作
c = tensor([4., 5., 1.]
c 只是张量 a 的元素之和,在张量 b 中出现两个 0 之间。
您可以使用torch.tensor_split
在 b 中的 0 索引上拆分张量,然后单独对它们求和:
例如:
group = torch.tensor_split(a, torch.where(b==0)[0])
# Output:
# (tensor([]),
# tensor([1., 1., 1., 1.]),
# tensor([1., 1., 1., 1., 1.]),
# tensor([1.]))
individual_sum = list(map(torch.sum, group)) # You can use loop/list comprehension etc
# Output
# [tensor(0.), tensor(4.), tensor(5.), tensor(1.)]
请注意,第一个 0 也被考虑,并在分割后产生一个空的张量。您可以在组合时删除它
torch.tensor(individual_sum[1:])
# Output
# tensor([4., 5., 1.])
本文收集自互联网,转载请注明来源。
如有侵权,请联系 [email protected] 删除。
我来说两句