python - 如何从索引向量 B 中计算从向量 A 中选择的值的平均值？

Question

我有一个值向量，例如:

import torch
v = torch.rand(6)

tensor([0.0811, 0.9658, 0.1901, 0.0872, 0.8895, 0.9647])

和一个索引来从v中选择值:

index = torch.tensor([0, 1, 0, 2, 0, 2])

我想生成一个向量 mean，它将计算 v 的平均值，这些值按 index 中的索引分组。

在这个例子中:

mean = torch.tensor([(0.0811 + 0.1901 + 0.8895) / 3, 0.9658, (0.0872 + 0.9647) / 2, 0, 0, 0])

tensor([0.3869, 0.9658, 0.5260, 0.0000, 0.0000, 0.0000])

Answer 1

一种可能的解决方案是结合使用 torch.bincount 和 Tensor.index_add():

v = torch.tensor([0.0811, 0.9658, 0.1901, 0.0872, 0.8895, 0.9647])
index = torch.tensor([0, 1, 0, 2, 0, 2])

bincount() 获取 index 中每个索引使用的总数:

bincount = torch.bincount(index, minlength=6)
# --> tensor([3, 1, 2, 0, 0, 0])

index_add() 按照 index 给定的顺序从 v 添加:

numerator = torch.zeros(6)
numerator = numerator.index_add(0, index, v)
# --> tensor([1.1607, 0.9658, 1.0520, 0.0000, 0.0000, 0.0000])

用 1.0 替换 bincount 中的零以防止被 0 除并从 int 转换为 float:

div = bincount.float()
div[bincount == 0] = 1.0
# --> tensor([3., 1., 2., 1., 1., 1.])

mean = num/div
# --> tensor([0.3869, 0.9658, 0.5260, 0.0000, 0.0000, 0.0000])