pytorch中的groupby聚合均值

时间:2019-05-15 17:18:36

标签: pytorch

我有一个二维张量:

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循环?

3 个答案:

答案 0 :(得分:1)

所有您需要做的是形成一个mxn矩阵(m = num类,n = num个样本),该矩阵将选择适当的权重并适当地缩放平均值。然后,您可以在新形成的矩阵和样本矩阵之间进行矩阵乘法。

给出标签,矩阵应该是(每行是一个类别编号,每个类别是一个样本编号及其权重):

public void createKey(Node root,String s,HashMap<String,String> hashMap){
    if(root==null)
        return;
    hashMap.put(root.symbol, s);
    createKey(root.left, s + "0",hashMap);
    createKey(root.right, s + "1",hashMap);
}

您可以形成以下形式:

[[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(4)] = 1
M = torch.nn.functional.normalize(M, p=1, dim=1)
torch.mm(M, samples)

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

答案 1 :(得分:0)

在此处{@ {3}}中将@ptrblck_de的答案重新发布

labels = labels.view(labels.size(0), 1).expand(-1, samples.size(1))

unique_labels, labels_count = labels.unique(dim=0, return_counts=True)

res = torch.zeros_like(unique_labels, dtype=torch.float).scatter_add_(0, labels, samples)
res = res / labels_count.float().unsqueeze(1)

答案 2 :(得分:0)

由于以前的解决方案不适用于稀疏组的情况(例如,并非所有组都在数据中),我做了一个 :)

def groupby_mean(value:torch.Tensor, labels:torch.LongTensor) -> (torch.Tensor, torch.LongTensor):
    """Group-wise average for (sparse) grouped tensors

    Args:
        value (torch.Tensor): values to average (# samples, latent dimension)
        labels (torch.LongTensor): labels for embedding parameters (# samples,)

    Returns: 
        result (torch.Tensor): (# unique labels, latent dimension)
        new_labels (torch.LongTensor): (# unique labels,)

    Examples:
        >>> samples = torch.Tensor([
                             [0.15, 0.15, 0.15],    #-> group / class 1
                             [0.2, 0.2, 0.2],    #-> group / class 3
                             [0.4, 0.4, 0.4],    #-> group / class 3
                             [0.0, 0.0, 0.0]     #-> group / class 0
                      ])
        >>> labels = torch.LongTensor([1, 5, 5, 0])
        >>> result, new_labels = groupby_mean(samples, labels)

        >>> result
        tensor([[0.0000, 0.0000, 0.0000],
            [0.1500, 0.1500, 0.1500],
            [0.3000, 0.3000, 0.3000]])

        >>> new_labels
        tensor([0, 1, 5])
    """
    uniques = labels.unique().tolist()
    labels = labels.tolist()

    key_val = {key: val for key, val in zip(uniques, range(len(uniques)))}
    val_key = {val: key for key, val in zip(uniques, range(len(uniques)))}

    labels = torch.LongTensor(list(map(key_val.get, labels)))

    labels = labels.view(labels.size(0), 1).expand(-1, value.size(1))

    unique_labels, labels_count = labels.unique(dim=0, return_counts=True)
    result = torch.zeros_like(unique_labels, dtype=torch.float).scatter_add_(0, labels, value)
    result = result / labels_count.float().unsqueeze(1)
    new_labels = torch.LongTensor(list(map(val_key.get, unique_labels[:, 0].tolist())))
    return result, new_labels
相关问题