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

我有一个值向量，例如:

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

请您参考如下方法:

一种可能的解决方案是结合使用 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])