基于另一个张量求和 Torch 张量

Question

我有两个张量，第一个包含浮点数，第二个包含 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.])

我想要一些矢量化（最好）的操作来接收两个张量和 returns

c = tensor([4., 5., 1.]

c只是张量a的元素之和，在张量b中出现的两个0之间。

Answer 1

您可以使用 torch.tensor_split 将张量拆分为 b 中 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.])
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.])

基于另一个张量求和 Torch 张量

Summing Torch tensor based on another tensor

python

vectorization

torch