Groupby aggregate mean in pytorch

elyase: <p>I have a 2D tensor:</p>
<pre><code class="lang-auto">    samples = torch.Tensor([
                     [0.1, 0.1],    #-&gt; group / class 1
                     [0.2, 0.2],    #-&gt; group / class 2
                     [0.4, 0.4],    #-&gt; group / class 2
                     [0.0, 0.0]     #-&gt; group / class 0
              ])
</code></pre>
<p>and a label for each sample corresponding to a class:</p>
<p><code>labels = torch.LongTensor([1, 2, 2, 0])</code></p>
<p>so <code>len(samples) == len(labels)</code>. Now I want to calculate the mean for each class / label. Because there are 3 classes (0, 1 and 2) the final vector should have dimension <code>[n_classes, samples.shape[1]]</code> So the expected solution should be:</p>
<pre><code class="lang-auto">    result == torch.Tensor([
                     [0.1, 0.1],
                     [0.3, 0.3], # -&gt; mean of [0.2, 0.2] and [0.4, 0.4]
                     [0.0, 0.0]
              ])
</code></pre>
<p><strong>Question</strong>: Can this be done in pure pytorch (i.e. no numpy so that I can autograd) and ideally without for loops?</p>

ptrblck: <p>You could use <code>scatter_add_</code> and <code>torch.unique</code> to get a similar result.<br>
However, the result tensor will be sorted according to the class index:</p>
<pre><code class="lang-python">samples = torch.Tensor([
                     [0.1, 0.1],    #-&gt; group / class 1
                     [0.2, 0.2],    #-&gt; group / class 2
                     [0.4, 0.4],    #-&gt; group / class 2
                     [0.0, 0.0]     #-&gt; group / class 0
              ])

labels = torch.LongTensor([1, 2, 2, 0])
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)
</code></pre>

elyase: <p>worked like a charm thanks a lot! I also posted on <a href="https://stackoverflow.com/a/56155805/1330293" rel="nofollow noopener">stackoverflow</a> and got the following alternative which I leave here for future reference:</p>
<pre><code class="lang-auto">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)
</code></pre>

ptrblck: <p>That’s also an interesting approach! Thanks for sharing.</p>

Wenbin_Zhu: <p>If your labels are sparse, like: [1, 2, 2, 2] (where 0 is missing), the first solution doesn’t work. The second one works but outputs the mean for label 0, here is my fix:</p>
<pre><code class="lang-auto">
def mean_by_label(samples, labels):
    ''' select mean(samples), count() from samples group by labels order by labels asc '''
    weight = torch.zeros(labels.max()+1, samples.shape[0]).to(samples.device) # L, N
    weight[labels, torch.arange(samples.shape[0])] = 1
    label_count = weight.sum(dim=1)
    weight = torch.nn.functional.normalize(weight, p=1, dim=1) # l1 normalization
    mean = torch.mm(weight, samples) # L, F
    index = torch.arange(mean.shape[0])[label_count &gt; 0]
    return mean[index], label_count[index]
</code></pre>

Eric1997: <p>Hi thanks for sharing your code! Do you know how to caculate the variance for each label (in a similar way)?</p>

Minseok: <p>As previous solutions do not work for the case of sparse groups (e.g., not all the groups are in the data), I made one <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>
<pre><code class="lang-auto">def groupby_mean(value:torch.Tensor, labels:torch.LongTensor) -&gt; (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:
        &gt;&gt;&gt; samples = torch.Tensor([
                             [0.15, 0.15, 0.15],    #-&gt; group / class 1
                             [0.2, 0.2, 0.2],    #-&gt; group / class 3
                             [0.4, 0.4, 0.4],    #-&gt; group / class 3
                             [0.0, 0.0, 0.0]     #-&gt; group / class 0
                      ])
        &gt;&gt;&gt; labels = torch.LongTensor([1, 5, 5, 0])
        &gt;&gt;&gt; result, new_labels = groupby_mean(samples, labels)
        
        &gt;&gt;&gt; result
        tensor([[0.0000, 0.0000, 0.0000],
            [0.1500, 0.1500, 0.1500],
            [0.3000, 0.3000, 0.3000]])
            
        &gt;&gt;&gt; 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
</code></pre>

ookey: <p>Hello,<br>
as mentionned, the case were one label never occurs fails. There are other solutions, but that output 0 for those missing labels. Any solution with the same format as <a class="mention" href="/u/ptrblck">@ptrblck</a> ?</p>