按一列中的特定值对行进行分组,然后在PyTorch中计算平均值

时间:2020-06-17 07:55:19

标签: python pytorch

样本张量:

tensor([[ 0.,  1.,  2.,  3.,  4.,  5.],  # class1
        [ 6.,  7.,  8.,  9., 10., 11.],  # class3
        [12., 13., 14., 15., 16., 17.],  # class2
        [18., 19., 20., 21., 22., 23.],  # class0
        [24., 25., 26., 27., 28., 29.].  # class1
])

预期结果:

tensor([[18., 19., 20., 21., 22., 23.], # class0
        [12., 13., 14., 15., 16., 17.], # class1
        [12., 13., 14., 15., 16., 17.], # class2
        [ 6.,  7.,  8.,  9., 10., 11.]. # class3
]) 

是否有一个纯PyTorch方法来实现?

1 个答案:

答案 0 :(得分:2)

您可以使用index_add根据类索引进行添加,然后除以每个标签的数量(使用unique计算得出):

# inputs
x = torch.arange(30.).view(5,6)  # sample tensor
c = c = torch.tensor([1, 3, 2, 0, 1], dtype=torch.long)  # class indices

# allocate space for output
result = torch.zeros((c.max() + 1, x.shape[1]), dtype=x.dtype)
# use index_add_ to sum up rows according to class
result.index_add_(0, c, x)
# use "unique" to count how many of each class
_, counts = torch.unique(c, return_counts=True)
# divide the sum by the counts to get the average
result /= counts[:, None]

result符合预期:

Out[*]:
tensor([[18., 19., 20., 21., 22., 23.],
        [12., 13., 14., 15., 16., 17.],
        [12., 13., 14., 15., 16., 17.],
        [ 6.,  7.,  8.,  9., 10., 11.]])