pytorch中的groupby聚合平均值

Question

我有一个二维张量：

samples = torch.Tensor([
    [0.1, 0.1],    #-> group / class 1
    [0.2, 0.2],    #-> group / class 2
    [0.4, 0.4],    #-> group / class 2
    [0.0, 0.0]     #-> group / class 0
])

以及对应于一个类的每个样本的标签：

labels = torch.LongTensor([1, 2, 2, 0])

所以len(samples) == len(labels)。现在我想计算每个类别/标签的平均值。因为有 3 个类（0、1 和 2），所以最终向量应该具有维度[n_classes, samples.shape[1]]所以预期的解决方案应该是：

result == torch.Tensor([
    [0.1, 0.1],
    [0.3, 0.3], # -> mean of [0.2, 0.2] and [0.4, 0.4]
    [0.0, 0.0]
])

问题：如何在纯 pytorch 中完成此操作（即没有 numpy，以便我可以自动分级）并且理想情况下没有 for 循环？

Answer 1

您需要做的就是形成一个 mxn 矩阵（m=num 个类，n=num 个样本），它将选择适当的权重，并适当缩放平均值。然后，您可以在新形成的矩阵和样本矩阵之间执行矩阵乘法。

给定你的标签，你的矩阵应该是（每行是一个类号，每个类一个样本号及其权重）：

[[0.0000, 0.0000, 0.0000, 1.0000],
 [1.0000, 0.0000, 0.0000, 0.0000],
 [0.0000, 0.5000, 0.5000, 0.0000]]

您可以按如下方式形成：

M = torch.zeros(labels.max()+1, len(samples))
M[labels, torch.arange(len(samples)] = 1
M = torch.nn.functional.normalize(M, p=1, dim=1)
torch.mm(M, samples)

输出：

tensor([[0.0000, 0.0000],
        [0.1000, 0.1000],
        [0.3000, 0.3000]])

请注意，输出均值按类顺序正确排序。

为什么M[labels, torch.arange(len(samples))] = 1有效？

这是在标签和样本数量之间执行广播操作。本质上，我们为标签中的每个元素生成一个 2D 索引：第一个指定它属于 m 个类中的哪一个，第二个仅指定其索引位置（从 1 到 N）。另一种方法是 top 显式生成所有 2D 索引：

twoD_indices = []
for count, label in enumerate(labels):
  twoD_indices.append((label, count))