Question

我有一个值向量，例如：

import torch
v = torch.rand(6)

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

和一个index从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（）按照v给出的顺序从index添加：

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

用bincount中的1.0替换零以防止被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])

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

1 个答案: