Hi,
I am wondering is there a efficient way to perform group averaging in a tensor.
Take 1-D vector as a sample:
vector: [1, 2, 3, 4, 5, 6]
group_label:[1, 2, 1, 2, 3, 3]
results:[2, 3, 5.5 ]
The ideas is that values in a tensor with same group label will be averaged and produce a new tensor.
1 Like
You could use scatter_add
and a small hack to get the unique counts of your indices:
x = torch.arange(1, 7, dtype=torch.float)
idx = torch.tensor([0, 1, 0, 1, 2, 2])
idx_unique = idx.unique(sorted=True)
idx_unique_count = torch.stack([(idx==idx_u).sum() for idx_u in idx_unique])
res = torch.zeros(len(idx_unique)).scatter_add(0, idx, x)
res /= idx_unique_count.float()
I think it’s time to add a return_counts
option to torch.unqiue
.
1 Like
Thanks for the help.
Thank you very much Sir.
thank you for your question!