группировать строки по заданному c значению в одном столбце и вычислять среднее значение в PyTorch - PullRequest
Новая
       5

группировать строки по заданному c значению в одном столбце и вычислять среднее значение в PyTorch

2 голосов
/ 17 июня 2020

образец тензора: 

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 Ответ

2 голосов
/ 17 июня 2020

Вы можете сложить в соответствии с индексом класса, используя 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.]])
...