Блог

У меня есть два тензора, где первый содержит поплавки, а второй содержит 0 и 1. Я хочу суммировать первый тензор на основе второго тензора. Более конкретно, я хочу подвести итог между появлением двух нулей. Например, рассмотрим

 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.

вы можете использовать torch.tensor_split , чтобы разделить свой тензор на индексы 0 в b, а затем суммировать их по отдельности:

Напр.:

 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.)]
 

Обратите внимание, что 1-й 0 также учитывается и приводит к пустому тензору после разделения. Вы можете удалить это во время объединения

 torch.tensor(individual_sum[1:])
# Output
# tensor([4., 5., 1.])